From dfc0d80944158265412f269e2b19791a0348d56b Mon Sep 17 00:00:00 2001 From: Hajime Kawahara Date: Tue, 28 Jan 2025 16:00:43 +0900 Subject: [PATCH 1/2] added test --- src/exojax/spec/opart.py | 81 ++++++++++++++++++- .../opart/opart_reflection_emis_test.py | 76 +++++++++++++++++ .../opart/opart_reflection_test.py | 9 +-- 3 files changed, 157 insertions(+), 9 deletions(-) create mode 100644 tests/integration/unittests_long/opart/opart_reflection_emis_test.py diff --git a/src/exojax/spec/opart.py b/src/exojax/spec/opart.py index 12de0cbc..6704e675 100644 --- a/src/exojax/spec/opart.py +++ b/src/exojax/spec/opart.py @@ -146,10 +146,7 @@ def update_layer(self, carry_rs, params): """ Rphat_prev, Sphat_prev = carry_rs - # - no source term - # temparature = params[0] - # source_vector = piB(temparature, self.nu_grid) - # ------------------------------------------------- + # no source term source_vector = jnp.zeros_like(self.nu_grid) dtau, single_scattering_albedo, asymmetric_parameter = self.opalayer(params) trans_coeff_i, scat_coeff_i, pihatB_i, _, _, _ = setrt_toonhm( @@ -190,6 +187,82 @@ def __call__( def run(self, opalayer, layer_params, flbl): return self(opalayer, layer_params, flbl) +class OpartReflectEmis(ArtCommon): + """Opart verision of ArtReflectEmis. + + This class computes the outgoing flux of the atmosphere with reflection, with emission from atmospheric layers. + Radiative transfer scheme: flux-based two-stream method, using flux-adding treatment, Toon-type hemispheric mean approximation + + """ + + def __init__(self, opalayer, pressure_top=1.0e-8, pressure_btm=1.0e2, nlayer=100): + """Initialization of OpartReflectPure + + Args: + opalayer (class): user defined class, needs to define self.nu_grid + pressure_top (float, optional): top pressure in bar. Defaults to 1.0e-8. + pressure_btm (float, optional): bottom pressure in bar. Defaults to 1.0e2. + nlayer (int, optional): the number of the atmospheric layers. Defaults to 100. + """ + super().__init__(pressure_top, pressure_btm, nlayer, opalayer.nu_grid) + self.opalayer = opalayer + self.nu_grid = self.opalayer.nu_grid + + def update_layer(self, carry_rs, params): + """updates the layer opacity and effective reflectivity (Rphat) and source (Sphat) + + Args: + carry_rs (list): carry for the Rphat and Sphat + params (list): layer parameters for this layer + + Returns: + list: updated carry_rs + """ + Rphat_prev, Sphat_prev = carry_rs + + # blackbody source term in the layers + temparature = params[0] + source_vector = piB(temparature, self.nu_grid) + # ------------------------------------------------- + dtau, single_scattering_albedo, asymmetric_parameter = self.opalayer(params) + trans_coeff_i, scat_coeff_i, pihatB_i, _, _, _ = setrt_toonhm( + dtau, single_scattering_albedo, asymmetric_parameter, source_vector + ) + denom = 1.0 - scat_coeff_i * Rphat_prev + Sphat_each = ( + pihatB_i + trans_coeff_i * (Sphat_prev + pihatB_i * Rphat_prev) / denom + ) + Rphat_each = scat_coeff_i + trans_coeff_i**2 * Rphat_prev / denom + carry_rs = [Rphat_each, Sphat_each] + return carry_rs + + def __call__( + self, + layer_params, + layer_update_function, + source_bottom, + refectivity_bottom, + incoming_flux, + ): + """computes outgoing flux + + Args: + layer_params (list): user defined layer parameters, layer_params[0] should be temperature array + layer_update_function (method): + source_bottom (array): source at the bottom [Nnus] + relfectivity_bottom (array): reflectivity at the bottom [Nnus] + incoming_flux (array): incoming flux [Nnus] + + Returns: + array: flux [Nnus] + """ + rs_bottom = [refectivity_bottom, source_bottom] + rs, _ = scan(layer_update_function, rs_bottom, layer_params) + return rs[0] * incoming_flux + rs[1] + + def run(self, opalayer, layer_params, flbl): + return self(opalayer, layer_params, flbl) + if __name__ == "__main__": diff --git a/tests/integration/unittests_long/opart/opart_reflection_emis_test.py b/tests/integration/unittests_long/opart/opart_reflection_emis_test.py new file mode 100644 index 00000000..6f8d1831 --- /dev/null +++ b/tests/integration/unittests_long/opart/opart_reflection_emis_test.py @@ -0,0 +1,76 @@ +"""checks the forward model of the opart spectrum +""" +import pytest +import numpy as np +import jax.numpy as jnp +from exojax.test.data import TESTDATA_CO_EXOMOL_PREMODIT_REFLECTION_REF +from exojax.test.emulate_mdb import mock_wavenumber_grid +from exojax.test.emulate_mdb import mock_mdbExomol +from exojax.spec.opacalc import OpaPremodit +from exojax.spec.opart import OpartReflectEmis +from exojax.spec.layeropacity import single_layer_optical_depth + +from jax import config +config.update("jax_enable_x64", True) + +def test_forward_reflection_emis_opart(): + class OpaLayer: + # user defined class, needs to define self.nugrid + def __init__(self): + self.nu_grid, self.wav, self.resolution = mock_wavenumber_grid() + self.gravity = 2478.57 + self.mdb_co = mock_mdbExomol() + + self.opa_co = OpaPremodit( + self.mdb_co, self.nu_grid, auto_trange=[400.0, 1500.0] + ) + + def __call__(self, params): + temperature, pressure, dP, mixing_ratio = params + xsv_co = self.opa_co.xsvector(temperature, pressure) + dtau_co = single_layer_optical_depth( + dP, xsv_co, mixing_ratio, self.mdb_co.molmass, self.gravity + ) + single_scattering_albedo = jnp.ones_like(dtau_co) * 0.3 + asymmetric_parameter = jnp.ones_like(dtau_co) * 0.01 + + return dtau_co, single_scattering_albedo, asymmetric_parameter + + opalayer = OpaLayer() + opart = OpartReflectEmis( + opalayer, pressure_top=1.0e-6, pressure_btm=1.0e0, nlayer=200 + ) + + def layer_update_function(carry, params): + carry = opart.update_layer(carry, params) + return carry, None + + temperature = opart.powerlaw_temperature(1300.0, 0.1) + mixing_ratio = opart.constant_mmr_profile(0.0003) + layer_params = [temperature, opart.pressure, opart.dParr, mixing_ratio] + + albedo = 1.0 + constant_incoming_flux = 1.e6 + incoming_flux = constant_incoming_flux*jnp.ones_like(opalayer.nu_grid) + reflectivity_surface = albedo * jnp.ones_like(opalayer.nu_grid) + + constant_surface_flux = 1.e5 + source_bottom = constant_surface_flux*jnp.ones_like(opalayer.nu_grid) + + flux = opart( + layer_params, layer_update_function, source_bottom, reflectivity_surface, incoming_flux + ) + ref = 1764352.3124300546 #2021 1/28 + print(np.mean(flux)) + assert np.mean(flux) == pytest.approx(ref) + plot = False + if plot: + import matplotlib.pyplot as plt + fig = plt.figure() + ax = fig.add_subplot(111) + plt.plot(opalayer.nu_grid, flux) + plt.savefig("forward_opart_reflect_emis.png") + + +if __name__ == "__main__": + test_forward_reflection_emis_opart() diff --git a/tests/integration/unittests_long/opart/opart_reflection_test.py b/tests/integration/unittests_long/opart/opart_reflection_test.py index 3b851e26..44664c25 100644 --- a/tests/integration/unittests_long/opart/opart_reflection_test.py +++ b/tests/integration/unittests_long/opart/opart_reflection_test.py @@ -1,6 +1,6 @@ """checks the forward model of the opart spectrum """ -import pkg_resources +from importlib.resources import files import pandas as pd import numpy as np import jax.numpy as jnp @@ -53,14 +53,13 @@ def layer_update_function(carry, params): albedo = 1.0 incoming_flux = jnp.ones_like(opalayer.nu_grid) reflectivity_surface = albedo * jnp.ones_like(opalayer.nu_grid) - source_bottom = jnp.zeros_like(opalayer.nu_grid) - + flux = opart( layer_params, layer_update_function, reflectivity_surface, incoming_flux ) - filename = pkg_resources.resource_filename('exojax', - 'data/testdata/' + TESTDATA_CO_EXOMOL_PREMODIT_REFLECTION_REF) + filename = files('exojax').joinpath('data/testdata/' + TESTDATA_CO_EXOMOL_PREMODIT_REFLECTION_REF) + dat = pd.read_csv(filename, delimiter=",", names=("nus", "flux")) residual = np.abs(flux / dat["flux"].values - 1.0) From 62ecdf81ea8a5adaff4021ffa5f0e96a44731e37 Mon Sep 17 00:00:00 2001 From: Hajime Kawahara Date: Tue, 28 Jan 2025 16:27:29 +0900 Subject: [PATCH 2/2] EmisScat (no surface) --- src/exojax/spec/opart.py | 94 +++++++++++++++++-- .../opart/opart_emis_scat_test.py | 65 +++++++++++++ 2 files changed, 150 insertions(+), 9 deletions(-) create mode 100644 tests/integration/unittests_long/opart/opart_emis_scat_test.py diff --git a/src/exojax/spec/opart.py b/src/exojax/spec/opart.py index 6704e675..6f7405fc 100644 --- a/src/exojax/spec/opart.py +++ b/src/exojax/spec/opart.py @@ -118,7 +118,7 @@ class OpartReflectPure(ArtCommon): This class computes the outgoing flux of the atmosphere with reflection, no emission from atmospheric layers nor surface. Radiative transfer scheme: flux-based two-stream method, using flux-adding treatment, Toon-type hemispheric mean approximation - + """ def __init__(self, opalayer, pressure_top=1.0e-8, pressure_btm=1.0e2, nlayer=100): @@ -164,7 +164,7 @@ def __call__( self, layer_params, layer_update_function, - refectivity_bottom, + reflectivity_bottom, incoming_flux, ): """computes outgoing flux @@ -180,19 +180,20 @@ def __call__( """ # rs_bottom = (refectivity_bottom, source_bottom) source_bottom = jnp.zeros_like(self.nu_grid) - rs_bottom = [refectivity_bottom, source_bottom] + rs_bottom = [reflectivity_bottom, source_bottom] rs, _ = scan(layer_update_function, rs_bottom, layer_params) return rs[0] * incoming_flux + rs[1] def run(self, opalayer, layer_params, flbl): return self(opalayer, layer_params, flbl) + class OpartReflectEmis(ArtCommon): """Opart verision of ArtReflectEmis. This class computes the outgoing flux of the atmosphere with reflection, with emission from atmospheric layers. Radiative transfer scheme: flux-based two-stream method, using flux-adding treatment, Toon-type hemispheric mean approximation - + """ def __init__(self, opalayer, pressure_top=1.0e-8, pressure_btm=1.0e2, nlayer=100): @@ -220,7 +221,7 @@ def update_layer(self, carry_rs, params): """ Rphat_prev, Sphat_prev = carry_rs - # blackbody source term in the layers + # blackbody source term in the layers temparature = params[0] source_vector = piB(temparature, self.nu_grid) # ------------------------------------------------- @@ -241,7 +242,7 @@ def __call__( layer_params, layer_update_function, source_bottom, - refectivity_bottom, + reflectivity_bottom, incoming_flux, ): """computes outgoing flux @@ -250,13 +251,13 @@ def __call__( layer_params (list): user defined layer parameters, layer_params[0] should be temperature array layer_update_function (method): source_bottom (array): source at the bottom [Nnus] - relfectivity_bottom (array): reflectivity at the bottom [Nnus] + reflectivity_bottom (array): reflectivity at the bottom [Nnus] incoming_flux (array): incoming flux [Nnus] - + Returns: array: flux [Nnus] """ - rs_bottom = [refectivity_bottom, source_bottom] + rs_bottom = [reflectivity_bottom, source_bottom] rs, _ = scan(layer_update_function, rs_bottom, layer_params) return rs[0] * incoming_flux + rs[1] @@ -264,6 +265,81 @@ def run(self, opalayer, layer_params, flbl): return self(opalayer, layer_params, flbl) +class OpartEmisScat(ArtCommon): + """Opart verision of ArtEmisScat. + + This class computes the outgoing emission flux of the atmosphere with scattering in the atmospheric layers. + Radiative transfer scheme: flux-based two-stream method, using flux-adding treatment, Toon-type hemispheric mean approximation + + """ + + def __init__(self, opalayer, pressure_top=1.0e-8, pressure_btm=1.0e2, nlayer=100): + """Initialization of OpartReflectPure + + Args: + opalayer (class): user defined class, needs to define self.nu_grid + pressure_top (float, optional): top pressure in bar. Defaults to 1.0e-8. + pressure_btm (float, optional): bottom pressure in bar. Defaults to 1.0e2. + nlayer (int, optional): the number of the atmospheric layers. Defaults to 100. + """ + super().__init__(pressure_top, pressure_btm, nlayer, opalayer.nu_grid) + self.opalayer = opalayer + self.nu_grid = self.opalayer.nu_grid + + def update_layer(self, carry_rs, params): + """updates the layer opacity and effective reflectivity (Rphat) and source (Sphat) + + Args: + carry_rs (list): carry for the Rphat and Sphat + params (list): layer parameters for this layer + + Returns: + list: updated carry_rs + """ + Rphat_prev, Sphat_prev = carry_rs + + # blackbody source term in the layers + temparature = params[0] + source_vector = piB(temparature, self.nu_grid) + # ------------------------------------------------- + dtau, single_scattering_albedo, asymmetric_parameter = self.opalayer(params) + trans_coeff_i, scat_coeff_i, pihatB_i, _, _, _ = setrt_toonhm( + dtau, single_scattering_albedo, asymmetric_parameter, source_vector + ) + denom = 1.0 - scat_coeff_i * Rphat_prev + Sphat_each = ( + pihatB_i + trans_coeff_i * (Sphat_prev + pihatB_i * Rphat_prev) / denom + ) + Rphat_each = scat_coeff_i + trans_coeff_i**2 * Rphat_prev / denom + carry_rs = [Rphat_each, Sphat_each] + return carry_rs + + def __call__( + self, + layer_params, + layer_update_function, + ): + """computes outgoing flux + + Args: + layer_params (list): user defined layer parameters, layer_params[0] should be temperature array + layer_update_function (method): + + Returns: + array: flux [Nnus] + """ + # no reflection at the bottom + reflectivity_bottom = jnp.zeros_like(self.nu_grid) + # no source term at the bottom + source_bottom = jnp.zeros_like(self.nu_grid) + rs_bottom = [reflectivity_bottom, source_bottom] + rs, _ = scan(layer_update_function, rs_bottom, layer_params) + return rs[1] + + def run(self, opalayer, layer_params, flbl): + return self(opalayer, layer_params, flbl) + + if __name__ == "__main__": from exojax.spec.opacalc import OpaPremodit diff --git a/tests/integration/unittests_long/opart/opart_emis_scat_test.py b/tests/integration/unittests_long/opart/opart_emis_scat_test.py new file mode 100644 index 00000000..278c5f62 --- /dev/null +++ b/tests/integration/unittests_long/opart/opart_emis_scat_test.py @@ -0,0 +1,65 @@ +"""checks the forward model of the opart spectrum +""" + +import pytest +import numpy as np +import jax.numpy as jnp +from exojax.test.emulate_mdb import mock_wavenumber_grid +from exojax.test.emulate_mdb import mock_mdbExomol +from exojax.spec.opacalc import OpaPremodit +from exojax.spec.opart import OpartEmisScat +from exojax.spec.layeropacity import single_layer_optical_depth + +from jax import config + +config.update("jax_enable_x64", True) + + +def test_forward_emis_scat_opart(): + class OpaLayer: + # user defined class, needs to define self.nugrid + def __init__(self): + self.nu_grid, self.wav, self.resolution = mock_wavenumber_grid() + self.gravity = 2478.57 + self.mdb_co = mock_mdbExomol() + + self.opa_co = OpaPremodit( + self.mdb_co, self.nu_grid, auto_trange=[400.0, 1500.0] + ) + + def __call__(self, params): + temperature, pressure, dP, mixing_ratio = params + xsv_co = self.opa_co.xsvector(temperature, pressure) + dtau_co = single_layer_optical_depth( + dP, xsv_co, mixing_ratio, self.mdb_co.molmass, self.gravity + ) + single_scattering_albedo = jnp.ones_like(dtau_co) * 0.3 + asymmetric_parameter = jnp.ones_like(dtau_co) * 0.01 + + return dtau_co, single_scattering_albedo, asymmetric_parameter + + opalayer = OpaLayer() + opart = OpartEmisScat(opalayer, pressure_top=1.0e-6, pressure_btm=1.0e0, nlayer=200) + + def layer_update_function(carry, params): + carry = opart.update_layer(carry, params) + return carry, None + + temperature = opart.powerlaw_temperature(1300.0, 0.1) + mixing_ratio = opart.constant_mmr_profile(0.0003) + layer_params = [temperature, opart.pressure, opart.dParr, mixing_ratio] + flux = opart(layer_params, layer_update_function) + print(np.mean(flux)) + ref = 515245.12625256577 # 1/28 2025 + assert np.mean(flux) == pytest.approx(ref) + plot = False + if plot: + import matplotlib.pyplot as plt + + fig = plt.figure() + plt.plot(opalayer.nu_grid, flux) + plt.savefig("forward_opart_emis_scat.png") + + +if __name__ == "__main__": + test_forward_emis_scat_opart()