diff --git a/examples/2024-CLi-juno-mwr/inspect_emcee_chain.py b/examples/2024-CLi-juno-mwr/inspect_emcee_chain.py new file mode 100644 index 00000000..a1ee3b63 --- /dev/null +++ b/examples/2024-CLi-juno-mwr/inspect_emcee_chain.py @@ -0,0 +1,58 @@ +# import emcee +import corner +# %config InlineBackend.figure_format = "retina" +import numpy as np +import matplotlib.pyplot as plt +from scipy.stats import norm +# reader = emcee.backends.HDFBackend("run_mcmc.h5") + +# tau = reader.get_autocorr_time() +# burnin = int(2 * np.max(tau)) +# thin = int(0.5 * np.min(tau)) +# samples = reader.get_chain(discard=burnin, flat=True, thin=thin) +# log_prob_samples = reader.get_log_prob(discard=burnin, flat=True, thin=thin) + +import h5py + +h5=h5py.File('run_mcmc_5000.h5', 'r') +chain = h5['mcmc']['chain'][180:] +h5.close() + +flattened_chain = chain.reshape(-1, 2) + +# Create labels for the parameters +labels = ["qNH3 [ppmv]","Temperature [K]"] + +# Create the corner plot +fig, ax = plt.subplots(2, 2, figsize=(10, 10)) +corner.hist2d(flattened_chain[:, 0], flattened_chain[:, 1], ax=ax[1, 0]) +ax[1, 0].set_xlabel("qNH3 [ppmv]") +ax[1, 0].set_ylabel("Temperature [K]") +ax[1, 0].set_xlim([0,700]) + +# Plot histograms for each parameter +x = np.linspace(1, 700, 700) +for i in range(2): + ax[i, i].hist(flattened_chain[:, i], bins=30, color='blue', alpha=0.7, density=True) + ax[i, i].set_xlabel(labels[i]) + ax[i, i].set_ylabel('PDF') + means=np.mean(flattened_chain[:, i]) + stdev=np.std(flattened_chain[:, i]) + ax[i, i].axvline(means, color='b', linestyle='--', label=f'posterior: ({means:6.2f}, {stdev:5.2f}') + + # Plot prior distribution + mean, stddev = [(300, 100), (169, 10)][i] + + prior = norm.pdf(x, mean, stddev) + ax[i, i].plot(x, prior, color='red', linestyle='--', label=f'Prior: norm({mean}, {stddev})') + ax[i, i].legend() + x = np.linspace(131, 200, 70) + +ax[0, 0].set_xlim([0,700]) + +ax[0, 1].axis('off') # Disable the axis +ax[0, 1].set_visible(False) # Hide the axis +# Show the plot +plt.tight_layout() +# Show the plot +plt.savefig("emcee_cornerplot_5000.png") \ No newline at end of file diff --git a/examples/2024-CLi-juno-mwr/inspect_emcee_step.py b/examples/2024-CLi-juno-mwr/inspect_emcee_step.py new file mode 100644 index 00000000..793b456a --- /dev/null +++ b/examples/2024-CLi-juno-mwr/inspect_emcee_step.py @@ -0,0 +1,41 @@ +# import emcee + +# %config InlineBackend.figure_format = "retina" +import numpy as np +import matplotlib.pyplot as plt +from scipy.stats import norm +# reader = emcee.backends.HDFBackend("run_mcmc.h5") + +# tau = reader.get_autocorr_time() +# burnin = int(2 * np.max(tau)) +# thin = int(0.5 * np.min(tau)) +# samples = reader.get_chain(discard=burnin, flat=True, thin=thin) +# log_prob_samples = reader.get_log_prob(discard=burnin, flat=True, thin=thin) + +import h5py + +h5=h5py.File('run_mcmc_5000.h5', 'r') +chain = h5['mcmc']['chain'][:] +h5.close() + +flattened_chain = chain.reshape(-1, 2) + +# Create labels for the parameters +labels = ["qNH3 [ppmv]","Temperature [K]"] + +# Create the corner plot +fig, ax = plt.subplots(2,1, figsize=(17, 10)) +for iw in range(5): + ax[0].plot(range(5000),chain[:,iw,0]) + +ax[0].set_ylabel("qNH3 [ppmv]") +ax[0].set_xlim([0,5000]) + +for iw in range(5): + ax[1].plot(range(5000),chain[:,iw,1]) +ax[1].set_ylabel("Temperature [K]") +ax[1].set_xlim([0,5000]) +ax[1].set_xlabel("step") +plt.tight_layout() +# Show the plot +plt.savefig("emcee_inspect_step_5000.png") \ No newline at end of file diff --git a/examples/2024-CLi-juno-mwr/juno_mwr.inp b/examples/2024-CLi-juno-mwr/juno_mwr.inp index e2bcdb51..63791840 100644 --- a/examples/2024-CLi-juno-mwr/juno_mwr.inp +++ b/examples/2024-CLi-juno-mwr/juno_mwr.inp @@ -65,7 +65,7 @@ PressureScaleHeight = 30.E3 gamma = 1.42 # gamma = C_p/C_v -grav_acc1 = -23.3 +grav_acc1 = -27.01 #-23.3 sfloor = 0. diff --git a/examples/2024-CLi-juno-mwr/jupiter_atmos.yaml b/examples/2024-CLi-juno-mwr/jupiter_atmos.yaml deleted file mode 100644 index 7d185052..00000000 --- a/examples/2024-CLi-juno-mwr/jupiter_atmos.yaml +++ /dev/null @@ -1,24 +0,0 @@ - -comment: Juno MWR configure file, physics, RT and inversion - -species: - vapor = [H2O, NH3] - cloud = [H2O(s), H2O(l), NH3(s), NH3(l)] - tracer = [e-, Na] - -thermodynamics: - dry: - Rd : 3587. # mu = 2.3175 g/mol - gammad : 1.42 # gamma = C_p/C_v - ## H2O - eps1 : [7.767, 7.767, 7.767] - beta1 : [0. , 24.845, 24.845] - rcp1 : [0.15, 0.33, 0.33] - Ttriple1 : 273.16 - Ptriple1 : 611.7 - ## NH3 - eps2 : [7.335, 7.335 , 7.335] - rcp2 : [0.078, 0.16, 0.16] - beta2 : [0. , 23.67, 23.67] - Ttriple2 : 195.4 - Ptriple2 : 6060. diff --git a/examples/2024-CLi-juno-mwr/run_juno_emcee.py b/examples/2024-CLi-juno-mwr/run_juno_emcee.py new file mode 100644 index 00000000..d3431300 --- /dev/null +++ b/examples/2024-CLi-juno-mwr/run_juno_emcee.py @@ -0,0 +1,210 @@ +#! /usr/bin/env python3 +import numpy as np +import emcee +import sys, os +import matplotlib.pyplot as plt +import h5py + +sys.path.append("../python") +sys.path.append(".") + +from canoe import def_species, load_configure +from canoe.snap import def_thermo +from canoe.athena import Mesh, ParameterInput, Outputs, MeshBlock +# from canoe.harp import radiation_band, radiation + +# forward operater +def run_RT_modify_atmos(mb: MeshBlock, + adlnTdlnP: float=0., + pmin: float=0. , + pmax: float =0. + )-> np.array: + # adlnTdlnP=0.0 ## set as insensitive + mb.modify_dlnTdlnP(adlnTdlnP, pmin, pmax) + # adlnNH3dlnP = 0#.25 + # mb.modify_dlnNH3dlnP(adlnNH3dlnP, pmin, pmax) + + # for k in range(mb.k_st, mb.k_ed + 1): + # for j in range(mb.j_st, mb.j_ed + 1): + rad=mb.get_rad() + rad.cal_radiance(mb, mb.k_st, mb.j_st) + + nb=rad.get_num_bands() + tb=np.array([0.]*4*nb) + + for ib in range(nb): + print(rad.get_band(ib)) + toa=rad.get_band(ib).get_toa()[0] + tb[ib*4:ib*4+4]=toa + return tb + +def set_atmos_run_RT(qNH3: float, # ppmv + T0: float=180. # Kelvin + ): + + mb.construct_atmosphere(pin,qNH3,T0) + rad=mb.get_rad() + rad.cal_radiance(mb, mb.k_st, mb.j_st) + + nb=rad.get_num_bands() + tb=np.array([0.]*4*nb) + + for ib in range(nb): + # print(rad.get_band(ib)) + toa=rad.get_band(ib).get_toa()[0] + tb[ib*4:ib*4+4]=toa + # print(tb[4:]) + return tb[4:] + +# Define likelihood function +def ln_likelihood(theta, observations, observation_errors): + nh3, temperature = theta + + simulations = set_atmos_run_RT(nh3,temperature) # Use your forward operator here + residuals = observations - simulations + print(simulations) + print(observations) + # print(residuals) + chi_squared = np.sum((residuals / observation_errors) ** 2) + # print(chi_squared) + return -0.5 * chi_squared + +# Define priors for NH3 and temperature +def ln_prior(theta): + nh3, temperature = theta + + nh3_mean = 300 # Mean value for NH3 + nh3_stddev = 100 # Standard deviation for NH3 + temperature_mean = 169 # Mean value for temperature + temperature_stddev = 10 # Standard deviation for temperature 0.5% + + ln_prior_nh3 = -0.5 * ((nh3 - nh3_mean) / nh3_stddev)**2 - np.log(nh3_stddev * np.sqrt(2 * np.pi)) + ln_prior_temperature = -0.5 * ((temperature - temperature_mean) / temperature_stddev)**2 - np.log(temperature_stddev * np.sqrt(2 * np.pi)) + + if (0 < nh3 < 1000):# and (100 < temperature < 200): + return ln_prior_nh3 + ln_prior_temperature + return -np.inf # return negative infinity if parameters are outside allowed range + +# Combine likelihood and prior to get posterior +def ln_posterior(theta, observations, observation_errors): + prior = ln_prior(theta) + if not np.isfinite(prior): + return -np.inf + return prior + ln_likelihood(theta, observations, observation_errors) + + +## main +if __name__ == "__main__": + + ## extract TB observations from ZZ fitting results + observations=np.zeros((20,)) + obs=np.zeros((24,)) + pj=51 + mu=np.cos(np.array([0.,15.,30.,45.])/180.*np.pi) + print(mu) + for ch in range(6): + tb_file=h5py.File(f"/nfs/nuke/chengcli/JUNOMWR/zzhang/PJ{pj:02d}_Freq{ch}.h5","r") + if ch==0: + c0=tb_file['ModelTypeupdate1_MultiPJ_Mode1/Iter1/c0'][-1] ## the north polar + c1=tb_file['ModelTypeupdate1_MultiPJ_Mode1/Iter1/c1'][-1] ## the north polar + c2=tb_file['ModelTypeupdate1_MultiPJ_Mode1/Iter1/c2'][-1] ## the north polar + else: + c0=tb_file['ModelTypeupdate1_MultiPJ_Mode3/Iter2/c0'][-1] ## the north polar + c1=tb_file['ModelTypeupdate1_MultiPJ_Mode3/Iter2/c1'][-1] ## the north polar + c2=tb_file['ModelTypeupdate1_MultiPJ_Mode3/Iter2/c2'][-1] ## the north polar + tb_file.close() + + # Xr=1.0 ## \mu >0.6 + obs[(ch)*4:(ch+1)*4]=c0-c1*5.*(1-mu)+c2/0.04*0.5*(mu-0.8)*(1-mu) + ## discard CH1 + observations=obs[4:] + print(observations) + + # [740.51932939 732.39178625 708.02076917 667.58562359 474.58510281 + # 469.42513666 454.00808555 428.59559118 338.13016122 335.65949356 + # 328.01197674 314.60534003 251.9730167 250.46642377 245.71888005 + # 237.15115289 194.47971955 193.67185714 191.10407859 186.40702401 + # 141.18445694 141.06723252 140.59821156 139.46693178] + + ## initialize Canoe + global pin + pin = ParameterInput() + pin.load_from_file("juno_mwr.inp") + + vapors = pin.get_string("species", "vapor").split(", ") + clouds = pin.get_string("species", "cloud").split(", ") + tracers = pin.get_string("species", "tracer").split(", ") + + def_species(vapors=vapors, clouds=clouds, tracers=tracers) + def_thermo(pin) + + config = load_configure("juno_mwr.yaml") + # print(pin.get_real("problem", "qH2O.ppmv")) + + mesh = Mesh(pin) + mesh.initialize(pin) + + global mb + mb = mesh.meshblock(0) + +## run MCMC + + # Generate synthetic observations and errors (replace with your real data) + # observations = np.random.normal(size=20) # Replace with your observations + observation_errors_stddev = 0.03 * observations ## error = 5% + + # Initialize walkers + n_walkers = 5 + n_dimensions = 2 # nh3, temperature + initial_guess = [200.0, 150.0] # Initial guess for NH3 and temperature + initial_guesses = [initial_guess + 10*np.random.randn(n_dimensions) for _ in range(n_walkers)] + + filename = "run_mcmc_5000.h5" + backend = emcee.backends.HDFBackend(filename) + backend.reset(n_walkers, n_dimensions) + + # Set up the sampler + sampler = emcee.EnsembleSampler(n_walkers, n_dimensions, ln_posterior, args=(observations, observation_errors_stddev), backend=backend) + + # Run MCMC + n_steps = 5000 + sampler.run_mcmc(initial_guesses, n_steps, progress=True) + + # Extract samples + samples = sampler.get_chain() + + # Compute mean and standard deviation of samples + mean_nh3 = np.mean(samples[:,:,0]) + std_nh3 = np.std(samples[:,:,0]) + mean_temperature = np.mean(samples[:,:,1]) + std_temperature = np.std(samples[:,:,1]) + + print("NH3: Mean =", mean_nh3, "Standard Deviation =", std_nh3) + print("Temperature: Mean =", mean_temperature, "Standard Deviation =", std_temperature) + + + # Plot convergence diagnostics + fig, axes = plt.subplots(2, figsize=(10, 7), sharex=True) + for iwalk in range(n_walkers): + axes[0].plot(range(n_steps),sampler.get_chain()[:,iwalk,0].T, alpha=0.4) + axes[0].set_ylabel("NH3") + for iwalk in range(n_walkers): + axes[1].plot(range(n_steps),sampler.get_chain()[:,iwalk,1].T, alpha=0.4) + axes[1].plot(sampler.get_chain()[:,:,1].T, alpha=0.4) + axes[1].set_ylabel("Temperature") + axes[1].set_xlabel("Step") + plt.savefig("MCMC-5000Steps.png") + + # Flatten samples for posterior distribution plotting + flat_samples = sampler.get_chain(discard=100, thin=15, flat=True) + + # Plot posterior distributions + labels = ["NH3", "Temperature"] + fig, axes = plt.subplots(2, figsize=(10, 7), sharex=True) + for i in range(2): + ax = axes[i] + ax.hist(flat_samples[:, i], bins=50, color="skyblue", histtype="stepfilled", alpha=0.7) + ax.set_title(labels[i]) + ax.grid(True) + # plt.tight_layout() + plt.savefig("posterior-distributions-5000.png") diff --git a/examples/2024-CLi-juno-mwr/run_juno_mwr.py b/examples/2024-CLi-juno-mwr/run_juno_mwr.py index 011cd009..b55bfca7 100644 --- a/examples/2024-CLi-juno-mwr/run_juno_mwr.py +++ b/examples/2024-CLi-juno-mwr/run_juno_mwr.py @@ -1,16 +1,33 @@ #! /usr/bin/env python3 import sys, os - +import numpy as np sys.path.append("../python") sys.path.append(".") from canoe import def_species, load_configure from canoe.snap import def_thermo -from canoe.athena import Mesh, ParameterInput, Outputs +from canoe.athena import Mesh, ParameterInput, Outputs,MeshBlock from typing import Tuple from canoe.harp import radiation_band, radiation -# from canoe.snap import air_parcel +def set_atmos_run_RT(mb: MeshBlock, + qNH3: float, # ppmv + T0: float # Kelvin + ): + + mb.construct_atmosphere(pin,qNH3,T0) + rad=mb.get_rad() + rad.cal_radiance(mb, mb.k_st, mb.j_st) + + nb=rad.get_num_bands() + tb=np.array([0.]*4*nb) + + for ib in range(nb): + # print(rad.get_band(ib)) + toa=rad.get_band(ib).get_toa()[0] + tb[ib*4:ib*4+4]=toa + return tb + if __name__ == "__main__": pin = ParameterInput() @@ -33,12 +50,30 @@ # out = Outputs(mesh, pin) # out.make_outputs(mesh, pin) + + # exit() mb = mesh.meshblock(0) + rad=mb.get_rad() + rad.cal_radiance(mb, mb.k_st, mb.j_st) + + nb=rad.get_num_bands() + tb=np.array([0.]*4*nb) + + for ib in range(nb): + band1=rad.get_band(ib) + toa=rad.get_band(ib).get_toa()[0] + tb[ib*4:ib*4+4]=toa + print(tb) + # print(len(mesh.meshblocks())) # for mb in mesh.meshblocks(): # print(mb.block_size.nx1) + tb=set_atmos_run_RT(mb,320,169) + print(tb) + + exit() adlnTdlnP = 0.2 ##k pmin = 100.0e5 pmax = 800.0e5 @@ -48,17 +83,21 @@ pmin = 100.0e5 pmax = 800.0e5 # mb.modify_dlnNH3dlnP(adlnNH3dlnP, pmin, pmax) + xNH3=200 + T0=190 + # mb.construct_atmosphere(pin,xNH3,T0) + # for k in range(mb.k_st, mb.k_ed + 1): # for j in range(mb.j_st, mb.j_ed + 1): # print(str(k) + " " + str(j)) - mb.get_rad().cal_radiance(mb, mb.k_st, mb.j_st) + # mb.get_rad().cal_radiance(mb, mb.k_st, mb.j_st) - pin.set_string("job", "problem_id", "juno_mwr_modify_Temp") + # pin.set_string("job", "problem_id", "juno_mwr_modify_Temp") # pin.set_string("job", "problem_id", "juno_mwr_modify_NH3") - print(pin.get_string("job", "problem_id")) + # print(pin.get_string("job", "problem_id")) - out = Outputs(mesh, pin) - out.make_outputs(mesh, pin) + # out = Outputs(mesh, pin) + # out.make_outputs(mesh, pin) # run rt again diff --git a/python/pyathena.cpp b/python/pyathena.cpp index 87426936..d9390fc3 100644 --- a/python/pyathena.cpp +++ b/python/pyathena.cpp @@ -16,6 +16,7 @@ // utils #include +#include namespace py = pybind11; @@ -194,6 +195,11 @@ void init_athena(py::module &parent) { pmax); }) + .def("construct_atmosphere", + [](MeshBlock &mesh_block, ParameterInput *pin, Real xNH3, Real T0) { + return construct_atmosphere(&mesh_block, pin, xNH3, T0); + }) + .def( "get_rad", [](MeshBlock &mesh_block) { return mesh_block.pimpl->prad; }, diff --git a/src/inversion/composition_inversion.cpp b/src/inversion/composition_inversion.cpp index 7b33e2ae..f7462a85 100644 --- a/src/inversion/composition_inversion.cpp +++ b/src/inversion/composition_inversion.cpp @@ -34,5 +34,30 @@ void CompositionInversion::UpdateModel(MeshBlock *pmb, Application::Logger app("inversion"); app->Log("Update Model"); - // int is = pblock_->is, ie = pblock_->ie; + auto air = AirParcel::gather_from_primitive(pmb, k, ju_, pmb->is); + air.ToMoleFraction(); + + int iter = 0, max_iter = 200; + while (iter++ < max_iter) { + // read in vapors + for (auto n : GetSpeciesIndex()) { + air.w[n] = par[n]; + } + + // stop at just above P0 + for (int i = is; i <= ie; ++i) { + pthermo->Extrapolate(&air, -dlnp / 2., "dry"); + if (air.w[IPR] < P0) break; + } + + // extrapolate down to where air is + pthermo->Extrapolate(&air, log(P0 / air.w[IPR]), "dry"); + + // make up for the difference + Ts += T0 - air.w[IDN]; + if (std::abs(T0 - air.w[IDN]) < 0.01) break; + + // app->Log("Iteration #", iter); + // app->Log("T", air.w[IDN]); + } } diff --git a/src/utils/construct_atmosphere.cpp b/src/utils/construct_atmosphere.cpp new file mode 100644 index 00000000..4f654dcd --- /dev/null +++ b/src/utils/construct_atmosphere.cpp @@ -0,0 +1,243 @@ +// C/C++ +#include +#include +#include +#include +#include + +// athena +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// application +#include +#include + +// canoe +#include +#include +#include +// #include +#include +#include + +// climath +#include + +// snap +#include +#include + +// utils +#include +#include +#include + +// astro +// #include +// #include + +// harp +// #include +// #include + +// tracer +#include + +// snap +#include + +// special includes +// #include "juno_mwr_specs.hpp" + +// set up an adiabatic atmosphere +void construct_atmosphere(MeshBlock *pmb, ParameterInput *pin, Real NH3ppmv, Real T0) { + Application::Logger app("main"); + // app->Log("ProblemGenerator: juno"); + + app->Log("NH3.ppmv", NH3ppmv); + app->Log("T0", T0); + + auto pmy_mesh= pmb->pmy_mesh; + auto pthermo = Thermodynamics::GetInstance(); + auto pcoord = pmb->pcoord; + auto pimpl = pmb->pimpl; + + int is = pmb->is; + int ie = pmb->ie; + int js = pmb->js, ks = pmb->ks; + int je = pmb->je, ke = pmb->ke; + ke=ks; + je=js; + // mesh limits + Real x1min = pmy_mesh->mesh_size.x1min; + Real x1max = pmy_mesh->mesh_size.x1max; + + Real H0 = pcoord->GetPressureScaleHeight(); + + Real P0 = pin->GetReal("mesh", "ReferencePressure"); + + Real Tmin = pin->GetReal("problem", "Tmin"); + // thermodynamic constants + Real gamma = pin->GetReal("hydro", "gamma"); + Real Rd = pthermo->GetRd(); + Real cp = gamma / (gamma - 1.) * Rd; + + // index + auto pindex = IndexMap::GetInstance(); + int iH2O = pindex->GetVaporId("H2O"); + int iNH3 = pindex->GetVaporId("NH3"); + + // app->Log("index of H2O", iH2O); + // app->Log("index of NH3", iNH3); + + // set up an adiabatic atmosphere + int max_iter = 200, iter = 0; + Real dlnp = pcoord->dx1f(is) / H0; + + AirParcel air(AirParcel::Type::MoleFrac); + + // estimate surface temperature and pressure + Real Ps = P0 * exp(-x1min / H0); + Real Ts = T0 * pow(Ps / P0, Rd / cp); + Real xH2O = pin->GetReal("problem", "qH2O.ppmv") / 1.E6; +// Real xNH3 = pin->GetReal("problem", "qNH3.ppmv") / 1.E6; + Real xNH3 = NH3ppmv/ 1.E6; + // app->Log("xH2O", xH2O); + // app->Log("xNH3", xNH3); + // app->Log("x1min", x1min); + // app->Log("P0", P0); + // app->Log("Rd", Rd); + // app->Log("gamma", gamma); + + Real rh_max_nh3 = pin->GetOrAddReal("problem", "rh_max.NH3", 1.); + + while (iter++ < max_iter) { + // read in vapors + air.w[iH2O] = xH2O; + air.w[iNH3] = xNH3; + air.w[IPR] = Ps; + air.w[IDN] = Ts; + + // stop at just above P0 + for (int i = is; i <= ie; ++i) { + pthermo->Extrapolate(&air, -dlnp / 2., "dry"); + if (air.w[IPR] < P0) break; + } + + // extrapolate down to where air is + pthermo->Extrapolate(&air, log(P0 / air.w[IPR]), "dry"); + + // make up for the difference + Ts += T0 - air.w[IDN]; + if (std::abs(T0 - air.w[IDN]) < 0.01) break; + + // app->Log("Iteration #", iter); + // app->Log("T", air.w[IDN]); + } + + if (iter > max_iter) { + throw RuntimeError("ProblemGenerator", "maximum iteration reached"); + } + + // construct atmosphere from bottom up + air.ToMoleFraction(); + for (int k = ks; k <= ke; ++k) + for (int j = js; j <= je; ++j) { + air.SetZero(); + air.w[iH2O] = xH2O; + air.w[iNH3] = xNH3; + air.w[IPR] = Ps; + air.w[IDN] = Ts; + + int i = is; + for (; i <= ie; ++i) { + // check relative humidity + Real rh = get_relative_humidity(air, iNH3); + air.w[iNH3] *= std::min(rh_max_nh3 / rh, 1.); + + AirParcelHelper::distribute_to_primitive(pmb, ks, js, i, air); + + pthermo->Extrapolate(&air, -dlnp, "dry"); + + if (air.w[IDN] < Tmin) break; + } + + // Replace adiabatic atmosphere with isothermal atmosphere if temperature + // is too low + pthermo->Extrapolate(&air, dlnp, "dry"); + for (; i <= ie; ++i) { + pthermo->Extrapolate(&air, -dlnp, "isothermal"); + AirParcelHelper::distribute_to_primitive(pmb, ks, js, i, air); + } + } + + + // set tracers, electron and Na + int ielec = pindex->GetTracerId("e-"); + int iNa = pindex->GetTracerId("Na"); + auto ptracer = pimpl->ptracer; + + Real xH2S = pin->GetReal("problem", "xH2S"); + + Real metallicity = pin->GetOrAddReal("problem", "metallicity", 0.); + + Real xNa = pin->GetReal("problem", "xNa"); + xNa *= pow(10., metallicity); + + Real xKCl = pin->GetReal("problem", "xKCl"); + xKCl *= pow(10., metallicity); + + Real xHe = pin->GetReal("problem", "xHe"); + + Real xCH4 = pin->GetReal("problem", "xCH4"); + + auto phydro=pmb->phydro; + + for (int k = ks; k <= ke; ++k) + for (int j = js; j <= je; ++j) + for (int i = is; i <= ie; ++i) { + Real temp = pthermo->GetTemp(pmb, k, j, i); + Real pH2S = xH2S * phydro->w(IPR, k, j, i); + Real pNa = xNa * phydro->w(IPR, k, j, i); + Real svp = sat_vapor_p_Na_H2S_Visscher(temp, pH2S); + pNa = std::min(svp, pNa); + + ptracer->u(iNa, k, j, i) = pNa / (Constants::kBoltz * temp); + ptracer->u(ielec, k, j, i) = saha_ionization_electron_density( + temp, ptracer->u(iNa, k, j, i), 5.14); + } + + auto peos=pmb->peos; + auto pfield=pmb->pfield; + auto pscalars=pmb->pscalars; + auto pbval=pmb->pbval; + + + // primitive to conserved conversion (hydro) + peos->PrimitiveToConserved(phydro->w, pfield->bcc, phydro->u, pcoord, is, ie, + js, je, ks, ke); + + // conserved to primitive conversion (tracer) + peos->PassiveScalarConservedToPrimitive(pscalars->s, phydro->u, pscalars->r, + pscalars->r, pcoord, is, ie, js, je, + ks, ke); + + // Microwave radiative transfer needs temperatures at cell interfaces, which + // are interpolated from cell centered hydrodynamic variables. Normally, the + // boundary conditions are taken care of internally. But, since we call + // radiative tranfer directly in pgen, we would need to update the boundary + // conditions manually. The following lines of code updates the boundary + // conditions. + phydro->hbvar.SwapHydroQuantity(phydro->w, HydroBoundaryQuantity::prim); + pbval->ApplyPhysicalBoundaries(0., 0., pbval->bvars_main_int); + + +}; diff --git a/src/utils/construct_atmosphere.hpp b/src/utils/construct_atmosphere.hpp new file mode 100644 index 00000000..f5b4ef11 --- /dev/null +++ b/src/utils/construct_atmosphere.hpp @@ -0,0 +1,15 @@ +#ifndef SRC_UTILS_CONSTRUCT_ATMOSPHERE_DRY_HPP_ +#define SRC_UTILS_CONSTRUCT_ATMOSPHERE_DRY_HPP_ + +// C/C++ +#include +#include +#include +#include +#include + + +// set up an adiabatic atmosphere +void construct_atmosphere(MeshBlock *pmb, ParameterInput *pin, Real xNH3, Real T0); + +#endif // SRC_UTILS_CONSTRUCT_ATMOSPHERE_HPP_ \ No newline at end of file