diff --git a/photochem/extensions/__init__.py b/photochem/extensions/__init__.py index e69de29..31d309f 100644 --- a/photochem/extensions/__init__.py +++ b/photochem/extensions/__init__.py @@ -0,0 +1 @@ +from .gasgiants import EvoAtmosphereGasGiant \ No newline at end of file diff --git a/photochem/extensions/gasgiants.py b/photochem/extensions/gasgiants.py index af57b14..b8c616e 100644 --- a/photochem/extensions/gasgiants.py +++ b/photochem/extensions/gasgiants.py @@ -15,6 +15,56 @@ ### Extension of EvoAtmosphere class for gas giants ### +class GasGiantData(): + + def __init__(self, planet_radius, planet_mass, P_ref, thermo_file): + + # Save several inputs + self.planet_radius = planet_radius + self.planet_mass = planet_mass + self.P_ref = P_ref + + # Metallicity calculator + self.m = Metallicity(thermo_file) + + # Parameters using during initialization + # The factor of pressure the atmosphere extends + # compared to predicted quench points of gases + self.BOA_pressure_factor = 5.0 + + # If True, then the guessed initial condition will used + # quenching relations as an initial guess + self.initial_cond_with_quenching = True + + # Parameters for determining steady state + self.TOA_pressure_avg = 1.0e-7*1e6 # mean TOA pressure (dynes/cm^2) + self.max_dT_tol = 5 # The permitted difference between T in photochem and desired T + self.max_dlog10edd_tol = 0.2 # The permitted difference between Kzz in photochem and desired Kzz + self.freq_update_PTKzz = 1000 # step frequency to update PTKzz profile. + self.max_total_step = 100_000 # Maximum total allowed steps before giving up + self.min_step_conv = 300 # Min internal steps considered before convergence is allowed + self.verbose = True # print information or not? + self.freq_print = 100 # Frequency in which to print + + # Values that will be needed later. All of these set + # in `initialize_to_climate_equilibrium_PT` + self.P_clima_grid = None # The climate grid + self.metallicity = None + self.CtoO = None + # Below for interpolation + self.log10P_interp = None + self.T_interp = None + self.log10edd_interp = None + self.P_desired = None + self.T_desired = None + self.Kzz_desired = None + # Index of climate grid that is bottom of photochemical grid + self.ind_b = None + # information needed during robust stepping + self.total_step_counter = None + self.nerrors = None + self.robust_stepper_initialized = None + class EvoAtmosphereGasGiant(EvoAtmosphere): def __init__(self, mechanism_file, stellar_flux_file, planet_mass, planet_radius, @@ -58,32 +108,11 @@ def __init__(self, mechanism_file, stellar_flux_file, planet_mass, planet_radius ff.name ) - # Save inputs that matter - self.planet_radius = planet_radius - self.planet_mass = planet_mass - self.P_ref = P_ref - - # Parameters using during initialization - # The factor of pressure the atmosphere extends - # compared to predicted quench points of gases - self.BOA_pressure_factor = 5.0 - # If True, then the guessed initial condition will used - # quenching relations as an initial guess - self.initial_cond_with_quenching = True - # For computing chemical equilibrium at a metallicity. if thermo_file is None: thermo_file = mechanism_file - self.m = Metallicity(thermo_file) - # Parameters for determining steady state - self.TOA_pressure_avg = 1.0e-7*1e6 # mean TOA pressure (dynes/cm^2) - self.max_dT_tol = 5 # The permitted difference between T in photochem and desired T - self.max_dlog10edd_tol = 0.2 # The permitted difference between Kzz in photochem and desired Kzz - self.freq_update_PTKzz = 1000 # step frequency to update PTKzz profile. - self.max_total_step = 100_000 # Maximum total allowed steps before giving up - self.min_step_conv = 300 # Min internal steps considered before convergence is allowed - self.verbose = True # print information or not? - self.freq_print = 100 # Frequency in which to print + # New data + self.gdat = GasGiantData(planet_radius, planet_mass, P_ref, thermo_file) # Values in photochem to adjust self.var.verbose = 0 @@ -94,25 +123,6 @@ def __init__(self, mechanism_file, stellar_flux_file, planet_mass, planet_radius self.var.conv_longdy = 0.01 # threshold relative change that determines convergence self.var.custom_binary_diffusion_fcn = custom_binary_diffusion_fcn - # Values that will be needed later. All of these set - # in `initialize_to_climate_equilibrium_PT` - self.P_clima_grid = None # The climate grid - self.metallicity = None - self.CtoO = None - # Below for interpolation - self.log10P_interp = None - self.T_interp = None - self.log10edd_interp = None - self.P_desired = None - self.T_desired = None - self.Kzz_desired = None - # Index of climate grid that is bottom of photochemical grid - self.ind_b = None - # information needed during robust stepping - self.total_step_counter = None - self.nerrors = None - self.robust_stepper_initialized = None - def initialize_to_climate_equilibrium_PT(self, P_in, T_in, Kzz_in, metallicity, CtoO, rainout_condensed_atoms=True): """Initialized the photochemical model to a climate model result that assumes chemical equilibrium at some metallicity and C/O ratio. @@ -133,24 +143,26 @@ def initialize_to_climate_equilibrium_PT(self, P_in, T_in, Kzz_in, metallicity, CtoO = 2 is twice the solar C/O ratio. """ + gdat = self.gdat + if P_in.shape[0] != T_in.shape[0]: raise Exception('Input P and T must have same shape') if P_in.shape[0] != Kzz_in.shape[0]: raise Exception('Input P and Kzz must have same shape') # Save inputs - self.P_clima_grid = P_in - self.metallicity = metallicity - self.CtoO = CtoO + gdat.P_clima_grid = P_in + gdat.metallicity = metallicity + gdat.CtoO = CtoO # Compute chemical equilibrium along the whole P-T profile - mix, mubar = self.m.composition(T_in, P_in, CtoO, metallicity, rainout_condensed_atoms) + mix, mubar = gdat.m.composition(T_in, P_in, CtoO, metallicity, rainout_condensed_atoms) - if self.TOA_pressure_avg*3 > P_in[-1]: + if gdat.TOA_pressure_avg*3 > P_in[-1]: raise Exception('The photochemical grid needs to extend above the climate grid') # Altitude of P-T grid - P1, T1, mubar1, z1 = compute_altitude_of_PT(P_in, self.P_ref, T_in, mubar, self.planet_radius, self.planet_mass, self.TOA_pressure_avg) + P1, T1, mubar1, z1 = compute_altitude_of_PT(P_in, gdat.P_ref, T_in, mubar, gdat.planet_radius, gdat.planet_mass, gdat.TOA_pressure_avg) # If needed, extrapolate Kzz and mixing ratios if P1.shape[0] != Kzz_in.shape[0]: Kzz1 = np.append(Kzz_in,Kzz_in[-1]) @@ -162,7 +174,7 @@ def initialize_to_climate_equilibrium_PT(self, P_in, T_in, Kzz_in, metallicity, mix1 = mix # The gravity - grav1 = gravity(self.planet_radius, self.planet_mass, z1) + grav1 = gravity(gdat.planet_radius, gdat.planet_mass, z1) # Next, we compute the quench levels quench_levels = determine_quench_levels(T1, P1, Kzz1, mubar1, grav1) @@ -170,7 +182,7 @@ def initialize_to_climate_equilibrium_PT(self, P_in, T_in, Kzz_in, metallicity, # If desired, this bit applies quenched initial conditions, and recomputes # the altitude profile for this new mubar. - if self.initial_cond_with_quenching: + if gdat.initial_cond_with_quenching: # Apply quenching to mixing ratios mix1['CH4'][quench_levels[0]:] = mix1['CH4'][quench_levels[0]] @@ -197,21 +209,21 @@ def initialize_to_climate_equilibrium_PT(self, P_in, T_in, Kzz_in, metallicity, mubar1[j] += mix1[sp][j]*self.dat.species_mass[i] # Update z1 to get a new altitude profile - P1, T1, mubar1, z1 = compute_altitude_of_PT(P1, self.P_ref, T1, mubar1, self.planet_radius, self.planet_mass, self.TOA_pressure_avg) + P1, T1, mubar1, z1 = compute_altitude_of_PT(P1, gdat.P_ref, T1, mubar1, gdat.planet_radius, gdat.planet_mass, gdat.TOA_pressure_avg) # Save P-T-Kzz for later interpolation and corrections - self.log10P_interp = np.log10(P1.copy()[::-1]) - self.T_interp = T1.copy()[::-1] - self.log10edd_interp = np.log10(Kzz1.copy()[::-1]) - self.P_desired = P1.copy() - self.T_desired = T1.copy() - self.Kzz_desired = Kzz1.copy() + gdat.log10P_interp = np.log10(P1.copy()[::-1]) + gdat.T_interp = T1.copy()[::-1] + gdat.log10edd_interp = np.log10(Kzz1.copy()[::-1]) + gdat.P_desired = P1.copy() + gdat.T_desired = T1.copy() + gdat.Kzz_desired = Kzz1.copy() # Bottom of photochemical model will be at a pressure a factor # larger than the predicted quench pressure. - if P1[ind]*self.BOA_pressure_factor > P1[0]: + if P1[ind]*gdat.BOA_pressure_factor > P1[0]: raise Exception('BOA in photochemical model wants to be deeper than BOA of climate model.') - self.ind_b = np.argmin(np.abs(P1 - P1[ind]*self.BOA_pressure_factor)) + gdat.ind_b = np.argmin(np.abs(P1 - P1[ind]*gdat.BOA_pressure_factor)) self._initialize_atmosphere(P1, T1, Kzz1, z1, mix1) @@ -230,11 +242,13 @@ def reinitialize_to_new_climate_PT(self, P_in, T_in, Kzz_in, mix): mix : dict Mixing ratios of all species in the atmosphere - """ + """ + + gdat = self.gdat - if self.P_clima_grid is None: + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') - if not np.all(np.isclose(self.P_clima_grid,P_in)): + if not np.all(np.isclose(gdat.P_clima_grid,P_in)): raise Exception('Input pressure grid does not match saved pressure grid') if P_in.shape[0] != T_in.shape[0]: raise Exception('Input P and T must have same shape') @@ -258,7 +272,7 @@ def reinitialize_to_new_climate_PT(self, P_in, T_in, Kzz_in, mix): mubar = mubar + mix[sp]*species_mass[ind] # Compute altitude of P-T grid - P1, T1, mubar1, z1 = compute_altitude_of_PT(P_in, self.P_ref, T_in, mubar, self.planet_radius, self.planet_mass, self.TOA_pressure_avg) + P1, T1, mubar1, z1 = compute_altitude_of_PT(P_in, gdat.P_ref, T_in, mubar, gdat.planet_radius, gdat.planet_mass, gdat.TOA_pressure_avg) # If needed, extrapolte Kzz and mixing ratios if P1.shape[0] != Kzz_in.shape[0]: Kzz1 = np.append(Kzz_in,Kzz_in[-1]) @@ -270,23 +284,25 @@ def reinitialize_to_new_climate_PT(self, P_in, T_in, Kzz_in, mix): mix1 = mix # Save P-T-Kzz for later interpolation and corrections - self.log10P_interp = np.log10(P1.copy()[::-1]) - self.T_interp = T1.copy()[::-1] - self.log10edd_interp = np.log10(Kzz1.copy()[::-1]) - self.P_desired = P1.copy() - self.T_desired = T1.copy() - self.Kzz_desired = Kzz1.copy() + gdat.log10P_interp = np.log10(P1.copy()[::-1]) + gdat.T_interp = T1.copy()[::-1] + gdat.log10edd_interp = np.log10(Kzz1.copy()[::-1]) + gdat.P_desired = P1.copy() + gdat.T_desired = T1.copy() + gdat.Kzz_desired = Kzz1.copy() self._initialize_atmosphere(P1, T1, Kzz1, z1, mix1) def _initialize_atmosphere(self, P1, T1, Kzz1, z1, mix1): "Little helper function preventing code duplication." + gdat = self.gdat + # Compute TOA index - ind_t = np.argmin(np.abs(P1 - self.TOA_pressure_avg)) + ind_t = np.argmin(np.abs(P1 - gdat.TOA_pressure_avg)) # Shift z profile so that zero is at photochem BOA - z1_p = z1 - z1[self.ind_b] + z1_p = z1 - z1[gdat.ind_b] # Calculate the photochemical grid z_top = z1_p[ind_t] @@ -308,7 +324,7 @@ def _initialize_atmosphere(self, P1, T1, Kzz1, z1, mix1): den_p = P_p/(k_boltz*T_p) # Compute new planet radius - planet_radius_new = self.planet_radius + z1[self.ind_b] + planet_radius_new = gdat.planet_radius + z1[gdat.ind_b] # Update photochemical model grid self.dat.planet_radius = planet_radius_new @@ -345,8 +361,11 @@ def return_atmosphere_climate_grid(self): ------- dict Contains temperature, Kzz, and mixing ratios. - """ - if self.P_clima_grid is None: + """ + + gdat = self.gdat + + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') # return full atmosphere @@ -354,8 +373,8 @@ def return_atmosphere_climate_grid(self): # Interpolate full atmosphere to clima grid sol = {} - sol['pressure'] = self.P_clima_grid.copy() - log10Pclima = np.log10(self.P_clima_grid[::-1]).copy() + sol['pressure'] = gdat.P_clima_grid.copy() + log10Pclima = np.log10(gdat.P_clima_grid[::-1]).copy() log10P = np.log10(out['pressure'][::-1]).copy() T = np.interp(log10Pclima, log10P, out['temperature'][::-1].copy()) @@ -386,9 +405,11 @@ def return_atmosphere(self, include_deep_atmosphere = True, equilibrium = False, ------- dict Contains temperature, Kzz, and mixing ratios. - """ + """ + + gdat = self.gdat - if self.P_clima_grid is None: + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') out = {} @@ -397,7 +418,7 @@ def return_atmosphere(self, include_deep_atmosphere = True, equilibrium = False, out['Kzz'] = self.var.edd species_names = self.dat.species_names[:(-2-self.dat.nsl)] if equilibrium: - mix, mubar = self.m.composition(out['temperature'], out['pressure'], self.CtoO, self.metallicity, rainout_condensed_atoms) + mix, mubar = gdat.m.composition(out['temperature'], out['pressure'], gdat.CtoO, gdat.metallicity, rainout_condensed_atoms) for key in mix: out[key] = mix[key] for key in species_names[:self.dat.np]: @@ -411,12 +432,12 @@ def return_atmosphere(self, include_deep_atmosphere = True, equilibrium = False, return out # Prepend the deeper atmosphere, which we will assume is at Equilibrium - inds = np.where(self.P_desired > self.wrk.pressure_hydro[0]) + inds = np.where(gdat.P_desired > self.wrk.pressure_hydro[0]) out1 = {} - out1['pressure'] = self.P_desired[inds] - out1['temperature'] = self.T_desired[inds] - out1['Kzz'] = self.Kzz_desired[inds] - mix, mubar = self.m.composition(out1['temperature'], out1['pressure'], self.CtoO, self.metallicity, rainout_condensed_atoms) + out1['pressure'] = gdat.P_desired[inds] + out1['temperature'] = gdat.T_desired[inds] + out1['Kzz'] = gdat.Kzz_desired[inds] + mix, mubar = gdat.m.composition(out1['temperature'], out1['pressure'], gdat.CtoO, gdat.metallicity, rainout_condensed_atoms) out['pressure'] = np.append(out1['pressure'],out['pressure']) out['temperature'] = np.append(out1['temperature'],out['temperature']) @@ -436,14 +457,15 @@ def initialize_robust_stepper(self, usol): ---------- usol : ndarray[double,dim=2] Input number densities - """ - if self.P_clima_grid is None: + """ + gdat = self.gdat + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') - self.total_step_counter = 0 - self.nerrors = 0 + gdat.total_step_counter = 0 + gdat.nerrors = 0 self.initialize_stepper(usol) - self.robust_stepper_initialized = True + gdat.robust_stepper_initialized = True def robust_step(self): """Takes a single robust integrator step @@ -455,11 +477,14 @@ def robust_step(self): then the algorithm things it is time to give up on reaching a steady state. If reached_steady_state then the algorithm has reached a steady state within tolerance. - """ - if self.P_clima_grid is None: + """ + + gdat = self.gdat + + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') - if not self.robust_stepper_initialized: + if not gdat.robust_stepper_initialized: raise Exception('This routine can only be called after `initialize_robust_stepper`') give_up = False @@ -468,15 +493,15 @@ def robust_step(self): for i in range(1): try: self.step() - self.total_step_counter += 1 + gdat.total_step_counter += 1 except PhotoException as e: # If there is an error, lets reinitialize, but get rid of any # negative numbers usol = np.clip(self.wrk.usol.copy(),a_min=1.0e-40,a_max=np.inf) self.initialize_stepper(usol) - self.nerrors += 1 + gdat.nerrors += 1 - if self.nerrors > 10: + if gdat.nerrors > 10: give_up = True break @@ -485,49 +510,49 @@ def robust_step(self): # Compute the max difference between the P-T profile in photochemical model # and the desired P-T profile - T_p = np.interp(np.log10(self.wrk.pressure_hydro.copy()[::-1]), self.log10P_interp, self.T_interp) + T_p = np.interp(np.log10(self.wrk.pressure_hydro.copy()[::-1]), gdat.log10P_interp, gdat.T_interp) T_p = T_p.copy()[::-1] max_dT = np.max(np.abs(T_p - self.var.temperature)) # Compute the max difference between the P-edd profile in photochemical model # and the desired P-edd profile - log10edd_p = np.interp(np.log10(self.wrk.pressure_hydro.copy()[::-1]), self.log10P_interp, self.log10edd_interp) + log10edd_p = np.interp(np.log10(self.wrk.pressure_hydro.copy()[::-1]), gdat.log10P_interp, gdat.log10edd_interp) log10edd_p = log10edd_p.copy()[::-1] max_dlog10edd = np.max(np.abs(log10edd_p - np.log10(self.var.edd))) # TOA pressure TOA_pressure = self.wrk.pressure_hydro[-1] - condition1 = converged and self.wrk.nsteps > self.min_step_conv or self.wrk.tn > self.var.equilibrium_time - condition2 = max_dT < self.max_dT_tol and max_dlog10edd < self.max_dlog10edd_tol and self.TOA_pressure_avg/3 < TOA_pressure < self.TOA_pressure_avg*3 + condition1 = converged and self.wrk.nsteps > gdat.min_step_conv or self.wrk.tn > self.var.equilibrium_time + condition2 = max_dT < gdat.max_dT_tol and max_dlog10edd < gdat.max_dlog10edd_tol and gdat.TOA_pressure_avg/3 < TOA_pressure < gdat.TOA_pressure_avg*3 if condition1 and condition2: - if self.verbose: + if gdat.verbose: print('nsteps = %i longdy = %.1e max_dT = %.1e max_dlog10edd = %.1e TOA_pressure = %.1e'% \ - (self.total_step_counter, self.wrk.longdy, max_dT, max_dlog10edd, TOA_pressure/1e6)) + (gdat.total_step_counter, self.wrk.longdy, max_dT, max_dlog10edd, TOA_pressure/1e6)) # success! reached_steady_state = True break - if not (self.wrk.nsteps % self.freq_update_PTKzz) or (condition1 and not condition2): + if not (self.wrk.nsteps % gdat.freq_update_PTKzz) or (condition1 and not condition2): # After ~1000 steps, lets update P,T, edd and vertical grid, if possible. try: - self.set_press_temp_edd(self.P_desired,self.T_desired,self.Kzz_desired,hydro_pressure=True) + self.set_press_temp_edd(gdat.P_desired,gdat.T_desired,gdat.Kzz_desired,hydro_pressure=True) except PhotoException: pass try: - self.update_vertical_grid(TOA_pressure=self.TOA_pressure_avg) + self.update_vertical_grid(TOA_pressure=gdat.TOA_pressure_avg) except PhotoException: pass self.initialize_stepper(self.wrk.usol) - if self.total_step_counter > self.max_total_step: + if gdat.total_step_counter > gdat.max_total_step: give_up = True break - if not (self.wrk.nsteps % self.freq_print) and self.verbose: + if not (self.wrk.nsteps % gdat.freq_print) and gdat.verbose: print('nsteps = %i longdy = %.1e max_dT = %.1e max_dlog10edd = %.1e TOA_pressure = %.1e'% \ - (self.total_step_counter, self.wrk.longdy, max_dT, max_dlog10edd, TOA_pressure/1e6)) + (gdat.total_step_counter, self.wrk.longdy, max_dT, max_dlog10edd, TOA_pressure/1e6)) return give_up, reached_steady_state @@ -556,20 +581,22 @@ def model_state_to_dict(self): state. This dictionary can be used as an input to "initialize_from_dict". """ - if self.P_clima_grid is None: + gdat = self.gdat + + if gdat.P_clima_grid is None: raise Exception('This routine can only be called after `initialize_to_climate_equilibrium_PT`') out = {} - out['P_clima_grid'] = self.P_clima_grid - out['metallicity'] = self.metallicity - out['CtoO'] = self.CtoO - out['log10P_interp'] = self.log10P_interp - out['T_interp'] = self.T_interp - out['log10edd_interp'] = self.log10edd_interp - out['P_desired'] = self.P_desired - out['T_desired'] = self.T_desired - out['Kzz_desired'] = self.Kzz_desired - out['ind_b'] = self.ind_b + out['P_clima_grid'] = gdat.P_clima_grid + out['metallicity'] = gdat.metallicity + out['CtoO'] = gdat.CtoO + out['log10P_interp'] = gdat.log10P_interp + out['T_interp'] = gdat.T_interp + out['log10edd_interp'] = gdat.log10edd_interp + out['P_desired'] = gdat.P_desired + out['T_desired'] = gdat.T_desired + out['Kzz_desired'] = gdat.Kzz_desired + out['ind_b'] = gdat.ind_b out['planet_radius_new'] = self.dat.planet_radius out['top_atmos'] = self.var.top_atmos out['temperature'] = self.var.temperature @@ -583,16 +610,18 @@ def initialize_from_dict(self, out): """Initializes the model from a dictionary created by the "model_state_to_dict" routine. """ - self.P_clima_grid = out['P_clima_grid'] - self.metallicity = out['metallicity'] - self.CtoO = out['CtoO'] - self.log10P_interp = out['log10P_interp'] - self.T_interp = out['T_interp'] - self.log10edd_interp = out['log10edd_interp'] - self.P_desired = out['P_desired'] - self.T_desired = out['T_desired'] - self.Kzz_desired = out['Kzz_desired'] - self.ind_b = out['ind_b'] + gdat = self.gdat + + gdat.P_clima_grid = out['P_clima_grid'] + gdat.metallicity = out['metallicity'] + gdat.CtoO = out['CtoO'] + gdat.log10P_interp = out['log10P_interp'] + gdat.T_interp = out['T_interp'] + gdat.log10edd_interp = out['log10edd_interp'] + gdat.P_desired = out['P_desired'] + gdat.T_desired = out['T_desired'] + gdat.Kzz_desired = out['Kzz_desired'] + gdat.ind_b = out['ind_b'] self.dat.planet_radius = out['planet_radius_new'] self.update_vertical_grid(TOA_alt=out['top_atmos']) self.set_temperature(out['temperature'])