From 1c4f0be2bffaeea8f8b27410f87ab432ec836d34 Mon Sep 17 00:00:00 2001 From: Torax team Date: Tue, 16 Apr 2024 05:18:16 -0700 Subject: [PATCH] Move Rmaj, Rmin, and B0 out of `Config` to the geometry arguments. This is one step in our improved config and time-dependent geometry changes. This actually removes time-dependent Rmaj, etc, but those will be added back in coming PRs. PiperOrigin-RevId: 625293182 --- torax/boundary_conditions.py | 2 +- torax/calc_coeffs.py | 6 +- torax/config.py | 6 -- torax/config_slice.py | 6 -- torax/geometry.py | 94 +++++++++++-------- torax/initial_states.py | 8 +- torax/physics.py | 19 ++-- torax/sim.py | 1 - torax/sources/bootstrap_current_source.py | 5 +- torax/sources/source_models.py | 10 +- torax/tests/physics.py | 4 +- .../test_data/test_iterbaseline_mockup.py | 6 +- .../tests/test_data/test_iterhybrid_mockup.py | 6 +- .../tests/test_data/test_iterhybrid_newton.py | 6 +- .../test_iterhybrid_predictor_corrector.py | 6 +- .../tests/test_data/test_iterhybrid_rampup.py | 6 +- torax/tests/test_lib/explicit_stepper.py | 1 - torax/tests/test_lib/torax_refs.py | 18 ++-- torax/transport_model/critical_gradient.py | 4 +- torax/transport_model/qlknn_wrapper.py | 10 +- 20 files changed, 102 insertions(+), 122 deletions(-) diff --git a/torax/boundary_conditions.py b/torax/boundary_conditions.py index e961ee4c..dea50b2c 100644 --- a/torax/boundary_conditions.py +++ b/torax/boundary_conditions.py @@ -54,7 +54,7 @@ def compute_boundary_conditions( # pylint: disable=invalid-name nGW = ( dynamic_config_slice.Ip - / (jnp.pi * dynamic_config_slice.Rmin**2) + / (jnp.pi * geo.Rmin**2) * 1e20 / dynamic_config_slice.nref ) diff --git a/torax/calc_coeffs.py b/torax/calc_coeffs.py index ab6a56bc..a99c40c6 100644 --- a/torax/calc_coeffs.py +++ b/torax/calc_coeffs.py @@ -320,12 +320,10 @@ def _calc_coeffs_full( source_models, implicit_source_profiles, geo, - dynamic_config_slice.Rmaj, ) + source_models_lib.sum_sources_psi( source_models, explicit_source_profiles, geo, - dynamic_config_slice.Rmaj, ) true_ne_face = core_profiles.ne.face_value() * dynamic_config_slice.nref @@ -343,7 +341,7 @@ def _calc_coeffs_full( * sigma * consts.mu0 / geo.J**2 - / dynamic_config_slice.Rmaj + / geo.Rmaj ) tic_psi = jnp.ones_like(toc_psi) toc_dens_el = geo.vpr @@ -517,7 +515,7 @@ def _calc_coeffs_full( # pylint: disable=invalid-name nGW = ( dynamic_config_slice.Ip - / (jnp.pi * dynamic_config_slice.Rmin**2) + / (jnp.pi * geo.Rmin**2) * 1e20 / dynamic_config_slice.nref ) diff --git a/torax/config.py b/torax/config.py index 8b1378da..16916c6b 100644 --- a/torax/config.py +++ b/torax/config.py @@ -193,10 +193,6 @@ class Config: """Configuration parameters for the `torax` module.""" # physical inputs - # major radius (R) in meters - Rmaj: TimeDependentField = 6.2 - # minor radius (a) in meters - Rmin: TimeDependentField = 2.0 # amu of main ion (if multiple isotope, make average) Ai: float = 2.5 # charge of main ion @@ -205,8 +201,6 @@ class Config: # Note that if Ip_from_parameters=False in geometry, then this Ip will be # overwritten by values from the geometry data Ip: TimeDependentField = 15.0 - # Toroidal magnetic field on axis [T] - B0: TimeDependentField = 5.3 # needed for qlknn and fusion power Zeff: TimeDependentField = 1.0 Zimp: TimeDependentField = 10.0 # impurity charge state assumed for dilution diff --git a/torax/config_slice.py b/torax/config_slice.py index 8a2b5071..39eb9b04 100644 --- a/torax/config_slice.py +++ b/torax/config_slice.py @@ -79,18 +79,12 @@ class DynamicConfigSlice: sources: Mapping[str, DynamicSourceConfigSlice] # physical inputs - # major radius (R) in meters - Rmaj: float - # minor radius (a) in meters - Rmin: float # amu of main ion (if multiple isotope, make average) Ai: float # charge of main ion Zi: float # total plasma current in MA Ip: float - # Toroidal magnetic field on axis [T] - B0: float # needed for qlknn and fusion power Zeff: float Zimp: float # impurity charge state assumed for dilution diff --git a/torax/geometry.py b/torax/geometry.py index 92da5e0d..ab402620 100644 --- a/torax/geometry.py +++ b/torax/geometry.py @@ -115,6 +115,8 @@ class Geometry: r_norm: jnp.ndarray r_face: jnp.ndarray r: jnp.ndarray + Rmaj: jnp.ndarray + Rmin: jnp.ndarray B0: jnp.ndarray volume: jnp.ndarray volume_face: jnp.ndarray @@ -172,12 +174,12 @@ class CHEASEGeometry(Geometry): delta_lower_face: jnp.ndarray -# pylint: enable=invalid-name - - def build_circular_geometry( config: config_lib.Config, kappa: float = 1.72, + Rmaj: float = 6.2, + Rmin: float = 2.0, + B0: float = 5.3, hires_fac: int = 4, ) -> CircularGeometry: """Constructs a CircularGeometry. @@ -192,6 +194,9 @@ def build_circular_geometry( config: General TORAX config. kappa: Elogination. Defaults to 1.72 for the ITER elongation, to approximately correct volume and area integral Jacobians. + Rmaj: major radius (R) in meters + Rmin: minor radius (a) in meters + B0: Toroidal magnetic field on axis [T] hires_fac: Grid refinement factor for poloidal flux <--> plasma current calculations. @@ -204,7 +209,7 @@ def build_circular_geometry( # Define mesh (Slab Uniform 1D with Jacobian = 1) dr_norm = jnp.array(1) / config.nr mesh = Grid1D.construct(nx=config.nr, dx=dr_norm) - rmax = jnp.array(config.Rmin) + rmax = jnp.array(Rmin) # helper variables for mesh cells and faces # r coordinate of faces r_face_norm = mesh.face_centers @@ -214,7 +219,8 @@ def build_circular_geometry( dr = dr_norm * rmax r_face = r_face_norm * rmax r = r_norm * rmax - B0 = jnp.array(config.B0) # pylint: disable=invalid-name + Rmaj = jnp.array(Rmaj) + B0 = jnp.array(B0) # assumed elongation profile on cell grid kappa_param = kappa @@ -222,18 +228,18 @@ def build_circular_geometry( # assumed elongation profile on cell grid kappa_face = 1 + r_face_norm * (kappa_param - 1) - volume = 2 * jnp.pi**2 * config.Rmaj * r**2 * kappa - volume_face = 2 * jnp.pi**2 * config.Rmaj * r_face**2 * kappa_face + volume = 2 * jnp.pi**2 * Rmaj * r**2 * kappa + volume_face = 2 * jnp.pi**2 * Rmaj * r_face**2 * kappa_face area = jnp.pi * r**2 * kappa area_face = jnp.pi * r_face**2 * kappa_face # V' for volume integrations vpr = ( - 4 * jnp.pi**2 * config.Rmaj * r * kappa + 4 * jnp.pi**2 * Rmaj * r * kappa + volume / kappa * (kappa_param - 1) / rmax ) vpr_face = ( - 4 * jnp.pi**2 * config.Rmaj * r_face * kappa_face + 4 * jnp.pi**2 * Rmaj * r_face * kappa_face + volume_face / kappa_face * (kappa_param - 1) / rmax ) # pylint: disable=invalid-name @@ -248,15 +254,15 @@ def build_circular_geometry( # uses <1/R^2> with circular geometry G2 = vpr / ( - 4 * jnp.pi**2 * config.Rmaj**2 * jnp.sqrt(1 - (r / config.Rmaj) ** 2) + 4 * jnp.pi**2 * Rmaj**2 * jnp.sqrt(1 - (r / Rmaj) ** 2) ) # generate G2_face by hand G2_outer_face = vpr_face[-1] / ( 4 * jnp.pi**2 - * config.Rmaj**2 - * jnp.sqrt(1 - (r_face[-1] / config.Rmaj) ** 2) + * Rmaj**2 + * jnp.sqrt(1 - (r_face[-1] / Rmaj) ** 2) ) G2_outer_face = jnp.expand_dims(G2_outer_face, 0) G2_face = jnp.concatenate( @@ -279,8 +285,8 @@ def build_circular_geometry( J = jnp.ones(len(r)) J_face = jnp.ones(len(r_face)) # simplified (constant) version of the F=B*R function - F = jnp.ones(len(r)) * config.Rmaj * B0 - F_face = jnp.ones(len(r_face)) * config.Rmaj * B0 + F = jnp.ones(len(r)) * Rmaj * B0 + F_face = jnp.ones(len(r_face)) * Rmaj * B0 # High resolution versions for j (plasma current) and psi (poloidal flux) # manipulations. Needed if psi is initialized from plasma current, which is @@ -288,21 +294,21 @@ def build_circular_geometry( r_hires_norm = jnp.linspace(0, 1, config.nr * hires_fac) r_hires = r_hires_norm * rmax - Rout = config.Rmaj + r - Rout_face = config.Rmaj + r_face + Rout = Rmaj + r + Rout_face = Rmaj + r_face - Rin = config.Rmaj - r - Rin_face = config.Rmaj - r_face + Rin = Rmaj - r + Rin_face = Rmaj - r_face # assumed elongation profile on hires grid kappa_hires = 1 + r_hires_norm * (kappa_param - 1) - volume_hires = 2 * jnp.pi**2 * config.Rmaj * r_hires**2 * kappa_hires + volume_hires = 2 * jnp.pi**2 * Rmaj * r_hires**2 * kappa_hires area_hires = jnp.pi * r_hires**2 * kappa_hires # V' for volume integrations on hires grid vpr_hires = ( - 4 * jnp.pi**2 * config.Rmaj * r_hires * kappa_hires + 4 * jnp.pi**2 * Rmaj * r_hires * kappa_hires + volume_hires / kappa_hires * (kappa_param - 1) / rmax ) # S' for area integrals on hires grid @@ -315,8 +321,8 @@ def build_circular_geometry( denom = ( 4 * jnp.pi**2 - * config.Rmaj**2 - * jnp.sqrt(1 - (r_hires / config.Rmaj) ** 2) + * Rmaj**2 + * jnp.sqrt(1 - (r_hires / Rmaj) ** 2) ) G2_hires = vpr_hires / denom @@ -345,6 +351,8 @@ def build_circular_geometry( r_norm=r_norm, r_face=r_face, r=r, + Rmaj=Rmaj, + Rmin=rmax, B0=B0, volume=volume, volume_face=volume_face, @@ -395,6 +403,9 @@ def build_chease_geometry( config: config_lib.Config, geometry_dir: str | None = None, geometry_file: str = 'ITER_hybrid_citrin_equil_cheasedata.mat2cols', + Rmaj: float = 6.2, + Rmin: float = 2.0, + B0: float = 5.3, hires_fac: int = 4, Ip_from_parameters: bool = True, ): @@ -414,6 +425,10 @@ def build_chease_geometry( geometry_dir is not provided, then it defaults to another dir. See implementation. geometry_file: CHEASE file name. + Rmaj: major radius (R) in meters. CHEASE geometries are normalized, so this + is used as an unnormalization factor. + Rmin: minor radius (a) in meters + B0: Toroidal magnetic field on axis [T]. hires_fac: Grid refinement factor for poloidal flux <--> plasma current calculations. Ip_from_parameters: If True, take Ip from parameter file and rescale psi. @@ -442,19 +457,20 @@ def build_chease_geometry( # grid. CHEASE variables are normalized. Need to unnormalize them with # reference values poloidal flux and CHEASE-internal-calculated plasma # current. - B0 = jnp.array(config.B0) # pylint: disable=invalid-name - psiunnormfactor = (config.Rmaj**2 * B0) * 2 * jnp.pi + Rmaj = jnp.array(Rmaj) + B0 = jnp.array(B0) + psiunnormfactor = (Rmaj**2 * B0) * 2 * jnp.pi psi_chease = chease_data['PSIchease=psi/2pi'] * psiunnormfactor Ip_chease = ( - chease_data['Ipprofile'] / constants.CONSTANTS.mu0 * config.Rmaj * B0 + chease_data['Ipprofile'] / constants.CONSTANTS.mu0 * Rmaj * B0 ) # toroidal flux coordinate - rho = chease_data['RHO_TOR=sqrt(Phi/pi/B0)'] * config.Rmaj + rho = chease_data['RHO_TOR=sqrt(Phi/pi/B0)'] * Rmaj rhon = chease_data['RHO_TOR_NORM'] # midplane radii - Rin_chease = chease_data['R_INBOARD'] * config.Rmaj - Rout_chease = chease_data['R_OUTBOARD'] * config.Rmaj + Rin_chease = chease_data['R_INBOARD'] * Rmaj + Rout_chease = chease_data['R_OUTBOARD'] * Rmaj # toroidal field flux function J_chease = chease_data['T=RBphi'] @@ -463,21 +479,21 @@ def build_chease_geometry( delta_lower_face_chease = chease_data['delta_bottom'] # flux surface integrals of various geometry quantities - C1 = chease_data['Int(Rdlp/|grad(psi)|)=Int(Jdchi)'] * config.Rmaj / B0 - C2 = chease_data['<1/R**2>'] * C1 / config.Rmaj**2 + C1 = chease_data['Int(Rdlp/|grad(psi)|)=Int(Jdchi)'] * Rmaj / B0 + C2 = chease_data['<1/R**2>'] * C1 / Rmaj**2 C3 = chease_data[''] * C1 * B0**2 - C4 = chease_data['<|grad(psi)|**2>'] * C1 * (B0 * config.Rmaj) ** 2 + C4 = chease_data['<|grad(psi)|**2>'] * C1 * (B0 * Rmaj) ** 2 # derived quantities for transport equations and transformations # <\nabla V> - g0_chease = 2 * jnp.pi * chease_data['<|grad(psi)|>'] * B0 * config.Rmaj * C1 + g0_chease = 2 * jnp.pi * chease_data['<|grad(psi)|>'] * B0 * Rmaj * C1 g1_chease = 4 * jnp.pi**2 * C1 * C4 # <(\nabla V)**2> g2_chease = 4 * jnp.pi**2 * C1 * C3 # <(\nabla V)**2 / R**2> g3_chease = C2[1:] / C1[1:] # <1/R**2> g3_chease = jnp.concatenate((jnp.array([1 / Rin_chease[0] ** 2]), g3_chease)) G2_chease = ( - config.Rmaj + Rmaj / (16 * jnp.pi**4) * J_chease[1:] * g2_chease[1:] @@ -519,8 +535,8 @@ def build_chease_geometry( Ip_scale_factor = 1 # volume, area, and dV/drho, dS/drho - volume_chease = chease_data['VOLUMEprofile'] * config.Rmaj**3 - area_chease = chease_data['areaprofile'] * config.Rmaj**2 + volume_chease = chease_data['VOLUMEprofile'] * Rmaj**3 + area_chease = chease_data['areaprofile'] * Rmaj**2 vpr_chease = math_utils.gradient(volume_chease, rho) spr_chease = math_utils.gradient(area_chease, rho) # gradient boundary approximation not appropriate here @@ -531,7 +547,7 @@ def build_chease_geometry( jtot_chease = ( 2 * jnp.pi - * config.Rmaj + * Rmaj * math_utils.gradient(Ip_chease, volume_chease) * Ip_scale_factor ) @@ -587,9 +603,9 @@ def build_chease_geometry( J_face = interp_func(r_face_norm) J = interp_func(r_norm) # simplified (constant) version of the F=B*R function - F = J * config.Rmaj * B0 + F = J * Rmaj * B0 # simplified (constant) version of the F=B*R function - F_face = J_face * config.Rmaj * B0 + F_face = J_face * Rmaj * B0 interp_func = lambda x: jnp.interp(x, rhon, psi_chease) psi_chease = interp_func(r_norm) @@ -660,6 +676,8 @@ def build_chease_geometry( r_norm=r_norm, r_face=r_face, r=r, + Rmaj=Rmaj, + Rmin=jnp.array(Rmin), B0=B0, volume=volume, volume_face=volume_face, diff --git a/torax/initial_states.py b/torax/initial_states.py index b4704f1f..0e8c9bc9 100644 --- a/torax/initial_states.py +++ b/torax/initial_states.py @@ -104,7 +104,7 @@ def _update_dens( # pylint: disable=invalid-name nGW = ( dynamic_config_slice.Ip - / (jnp.pi * dynamic_config_slice.Rmin**2) + / (jnp.pi * geo.Rmin**2) * 1e20 / dynamic_config_slice.nref ) @@ -373,7 +373,6 @@ def _calculate_currents_from_psi( jtot, jtot_face = physics.calc_jtot_from_psi( geo, psi, - dynamic_config_slice.Rmaj, ) bootstrap_profile = source_models.j_bootstrap.get_value( @@ -465,7 +464,7 @@ def _update_psi_from_j( scale = jnp.concatenate(( jnp.zeros((1,)), constants.CONSTANTS.mu0 - / (2 * jnp.pi * dynamic_config_slice.Rmaj * geo.G2_hires[1:]), + / (2 * jnp.pi * geo.Rmaj * geo.G2_hires[1:]), )) # dpsi_dr on the cell grid dpsi_dr_hires = scale * integrated @@ -563,7 +562,6 @@ def initial_core_profiles( geo=geo, jtot_face=currents.jtot_face, psi=psi, - Rmaj=dynamic_config_slice.Rmaj, q_correction_factor=dynamic_config_slice.q_correction_factor, ) s_face = physics.calc_s_from_psi(geo, psi) @@ -591,7 +589,6 @@ def initial_core_profiles( geo=geo, jtot_face=geo.jtot_face, psi=psi, - Rmaj=dynamic_config_slice.Rmaj, q_correction_factor=dynamic_config_slice.q_correction_factor, ) s_face = physics.calc_s_from_psi(geo, psi) @@ -642,7 +639,6 @@ def initial_core_profiles( core_profiles = physics.update_jtot_q_face_s_face( geo=geo, core_profiles=core_profiles, - Rmaj=dynamic_config_slice.Rmaj, q_correction_factor=dynamic_config_slice.q_correction_factor, ) diff --git a/torax/physics.py b/torax/physics.py index 8fac1eba..6ad43fa4 100644 --- a/torax/physics.py +++ b/torax/physics.py @@ -49,7 +49,6 @@ def get_main_ion_dilution_factor( def update_jtot_q_face_s_face( geo: Geometry, core_profiles: state.CoreProfiles, - Rmaj: float, q_correction_factor: float, ) -> state.CoreProfiles: """Updates jtot, jtot_face, q_face, and s_face.""" @@ -57,13 +56,11 @@ def update_jtot_q_face_s_face( jtot, jtot_face = calc_jtot_from_psi( geo, core_profiles.psi, - Rmaj, ) q_face, _ = calc_q_from_jtot_psi( geo, jtot_face, core_profiles.psi, - Rmaj, q_correction_factor, ) s_face = calc_s_from_psi( @@ -142,7 +139,6 @@ def calc_q_from_jtot_psi( geo: Geometry, jtot_face: jax.Array, psi: cell_variable.CellVariable, - Rmaj: float, q_correction_factor: float, ) -> tuple[jnp.ndarray, jnp.ndarray]: """Calculates q given jtot and psi. @@ -155,7 +151,6 @@ def calc_q_from_jtot_psi( geo: Magnetic geometry. jtot_face: Total toroidal current density on face grid. psi: Poloidal flux. - Rmaj: major radius (R) in meters. q_correction_factor: ad-hoc fix for non-physical circular geometry model such that q(r=a) = 3 for standard ITER parameters; @@ -170,7 +165,7 @@ def calc_q_from_jtot_psi( denom = jnp.abs(psi.face_grad()[1:] / geo.rmax) inv_iota = (2 * jnp.pi * geo.B0 * geo.r_face[1:]) / denom # use on-axis definition of q (Wesson 2004, Eq 3.48) - q0 = 2 * geo.B0 / (constants.CONSTANTS.mu0 * jtot_face[0] * Rmaj) + q0 = 2 * geo.B0 / (constants.CONSTANTS.mu0 * jtot_face[0] * geo.Rmaj) q0 = jnp.expand_dims(q0, 0) q_face = jnp.concatenate([q0, inv_iota]) q_face *= jnp.where( @@ -186,14 +181,12 @@ def calc_q_from_jtot_psi( def calc_jtot_from_psi( geo: Geometry, psi: cell_variable.CellVariable, - Rmaj: float, ) -> tuple[jnp.ndarray, jnp.ndarray]: """Calculates j from psi. Args: geo: Torus geometry. psi: Poloidal flux. - Rmaj: Major radius (R) in meters. Returns: jtot: total current density (Amps / m^2) on cell grid @@ -210,7 +203,9 @@ def calc_jtot_from_psi( # TODO flag to JAX team that jnp.gradient is not up-to-date, # and should allow for non-uniform grid differentiation - jtot_face = 2 * jnp.pi * Rmaj * math_utils.gradient(I_tot, geo.volume_face) + jtot_face = ( + 2 * jnp.pi * geo.Rmaj * math_utils.gradient(I_tot, geo.volume_face) + ) jtot = geometry.face_to_cell(jtot_face) @@ -259,7 +254,6 @@ def calc_nu_star( core_profiles: state.CoreProfiles, nref: float, Zeff: float, - Rmaj: float, coll_mult: float, ) -> jnp.ndarray: """Calculates nu star. @@ -269,7 +263,6 @@ def calc_nu_star( core_profiles: Core plasma profiles. nref: Reference value for normalization Zeff: Effective ion charge. - Rmaj: Major radius (R) in meters coll_mult: Collisionality multiplier in QLKNN for sensitivity testing. Returns: @@ -300,12 +293,12 @@ def calc_nu_star( ) # calculate bounce time - epsilon = geo.r_face / Rmaj + epsilon = geo.r_face / geo.Rmaj # to avoid divisions by zero epsilon = jnp.clip(epsilon, constants.CONSTANTS.eps) tau_bounce = ( core_profiles.q_face - * Rmaj + * geo.Rmaj / ( epsilon**1.5 * jnp.sqrt( diff --git a/torax/sim.py b/torax/sim.py index ec8d316e..6e5ad90f 100644 --- a/torax/sim.py +++ b/torax/sim.py @@ -389,7 +389,6 @@ def body_fun( output_state.core_profiles = physics.update_jtot_q_face_s_face( geo=geo, core_profiles=output_state.core_profiles, - Rmaj=dynamic_config_slice_t_plus_dt.Rmaj, q_correction_factor=dynamic_config_slice_t_plus_dt.q_correction_factor, ) diff --git a/torax/sources/bootstrap_current_source.py b/torax/sources/bootstrap_current_source.py index 00bf2e0e..d19364e1 100644 --- a/torax/sources/bootstrap_current_source.py +++ b/torax/sources/bootstrap_current_source.py @@ -108,13 +108,12 @@ def calc_neoclassical( geo, jtot_face, psi, - dynamic_config_slice.Rmaj, dynamic_config_slice.q_correction_factor, ) nuestar = ( 6.921e-18 * q_face - * dynamic_config_slice.Rmaj + * geo.Rmaj * true_ne_face * dynamic_config_slice.Zeff * lnLame @@ -126,7 +125,7 @@ def calc_neoclassical( nuistar = ( 4.9e-18 * q_face - * dynamic_config_slice.Rmaj + * geo.Rmaj * true_ni_face * dynamic_config_slice.Zeff**4 * lnLami diff --git a/torax/sources/source_models.py b/torax/sources/source_models.py index 6bc4b554..4ef26e0c 100644 --- a/torax/sources/source_models.py +++ b/torax/sources/source_models.py @@ -340,7 +340,6 @@ def sum_sources_psi( source_models: SourceModels, source_profile: source_profiles.SourceProfiles, geo: geometry.Geometry, - Rmaj: float, # pylint: disable=invalid-name ) -> jnp.ndarray: """Computes psi source values for sim.calc_coeffs.""" total = ( @@ -353,7 +352,7 @@ def sum_sources_psi( affected_core_profile=source_lib.AffectedCoreProfile.PSI.value, geo=geo, ) - denom = 2 * jnp.pi * Rmaj * geo.J**2 + denom = 2 * jnp.pi * geo.Rmaj * geo.J**2 scale_source = lambda src: -geo.vpr * src * constants.CONSTANTS.mu0 / denom return scale_source(total) @@ -435,7 +434,7 @@ def calc_and_sum_sources_psi( calculate_anyway=True, ) total += j_bootstrap_profiles.j_bootstrap - denom = 2 * jnp.pi * dynamic_config_slice.Rmaj * geo.J**2 + denom = 2 * jnp.pi * geo.Rmaj * geo.J**2 scale_source = lambda src: -geo.vpr * src * constants.CONSTANTS.mu0 / denom return scale_source(total), j_bootstrap_profiles.sigma @@ -486,7 +485,7 @@ def calc_psidot( * sigma * consts.mu0 / geo.J**2 - / dynamic_config_slice.Rmaj + / geo.Rmaj ) d_face_psi = geo.G2_face / geo.J_face / geo.rmax**2 @@ -510,12 +509,11 @@ def _ohmic_heat_model( jtot, _ = physics.calc_jtot_from_psi( geo, core_profiles.psi, - dynamic_config_slice.Rmaj, ) psidot = calc_psidot(source_models, dynamic_config_slice, geo, core_profiles) - pohm = jtot * psidot / (2 * jnp.pi * dynamic_config_slice.Rmaj) + pohm = jtot * psidot / (2 * jnp.pi * geo.Rmaj) return pohm diff --git a/torax/tests/physics.py b/torax/tests/physics.py index 399370a6..35a1fc2e 100644 --- a/torax/tests/physics.py +++ b/torax/tests/physics.py @@ -53,7 +53,6 @@ def test_calc_q_from_psi( geo, jtot, references.psi, - dynamic_config_slice.Rmaj, dynamic_config_slice.q_correction_factor, ) @@ -73,7 +72,7 @@ def calc_q_from_psi(config, geo): q[1:] = 1 / iota[1:] # Change from PINT: we don't read jtot from `geo` q[0] = ( - 2 * geo.B0 / (consts.mu0 * jtot[0] * config.Rmaj) + 2 * geo.B0 / (consts.mu0 * jtot[0] * geo.Rmaj) ) # use on-axis definition of q (Wesson 2004, Eq 3.48) q *= config.q_correction_factor @@ -137,7 +136,6 @@ def test_calc_jtot_from_psi( j, _ = physics.calc_jtot_from_psi( references.geo, references.psi, - references.config.Rmaj, ) np.testing.assert_allclose(j, references.jtot) diff --git a/torax/tests/test_data/test_iterbaseline_mockup.py b/torax/tests/test_data/test_iterbaseline_mockup.py index e862b2f7..bfdfc649 100644 --- a/torax/tests/test_data/test_iterbaseline_mockup.py +++ b/torax/tests/test_data/test_iterbaseline_mockup.py @@ -30,10 +30,7 @@ def get_config() -> config_lib.Config: resistivity_mult=200, Ip=15, # total plasma current in MA # physical inputs - Rmaj=6.2, # major radius (R) in meters - Rmin=2.0, # minor radius (a) in meters Ai=2.5, # amu of main ion (if multiple isotope, make average) - B0=5.3, # Toroidal magnetic field on axis [T] Zeff=1.74, # needed for qlknn and fusion power # effective impurity charge state assumed for matching dilution=0.862 Zimp=6.3623, @@ -159,6 +156,9 @@ def get_geometry(config: config_lib.Config) -> geometry.Geometry: config, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, # major radius (R) in meters + Rmin=2.0, # minor radius (a) in meters + B0=5.3, # Toroidal magnetic field on axis [T] ) diff --git a/torax/tests/test_data/test_iterhybrid_mockup.py b/torax/tests/test_data/test_iterhybrid_mockup.py index 8bad430d..765e07f4 100644 --- a/torax/tests/test_data/test_iterhybrid_mockup.py +++ b/torax/tests/test_data/test_iterhybrid_mockup.py @@ -33,10 +33,7 @@ def get_config() -> config_lib.Config: resistivity_mult=200, Ip=10.5, # total plasma current in MA # physical inputs - Rmaj=6.2, # major radius (R) in meters - Rmin=2.0, # minor radius (a) in meters Ai=2.5, # amu of main ion (if multiple isotope, make average) - B0=5.3, # Toroidal magnetic field on axis [T] Zeff=1.6, # needed for qlknn and fusion power # effective impurity charge state assumed for matching dilution=0.862. Zimp=10, @@ -166,6 +163,9 @@ def get_geometry(config: config_lib.Config) -> geometry.Geometry: config, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, # major radius (R) in meters + Rmin=2.0, # minor radius (a) in meters + B0=5.3, # Toroidal magnetic field on axis [T] ) diff --git a/torax/tests/test_data/test_iterhybrid_newton.py b/torax/tests/test_data/test_iterhybrid_newton.py index 58b8d645..9e744e45 100644 --- a/torax/tests/test_data/test_iterhybrid_newton.py +++ b/torax/tests/test_data/test_iterhybrid_newton.py @@ -37,10 +37,7 @@ def get_config() -> config_lib.Config: resistivity_mult=1, Ip=10.5, # total plasma current in MA # physical inputs - Rmaj=6.2, # major radius (R) in meters - Rmin=2.0, # minor radius (a) in meters Ai=2.5, # amu of main ion (if multiple isotope, make average) - B0=5.3, # Toroidal magnetic field on axis [T] Zeff=1.6, # needed for qlknn and fusion power # effective impurity charge state assumed for matching dilution=0.862. Zimp=10, @@ -182,6 +179,9 @@ def get_geometry(config: config_lib.Config) -> geometry.Geometry: config, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, # major radius (R) in meters + Rmin=2.0, # minor radius (a) in meters + B0=5.3, # Toroidal magnetic field on axis [T] ) diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.py b/torax/tests/test_data/test_iterhybrid_predictor_corrector.py index db28c402..2bccce50 100644 --- a/torax/tests/test_data/test_iterhybrid_predictor_corrector.py +++ b/torax/tests/test_data/test_iterhybrid_predictor_corrector.py @@ -33,10 +33,7 @@ def get_config() -> config_lib.Config: resistivity_mult=200, Ip=10.5, # total plasma current in MA # physical inputs - Rmaj=6.2, # major radius (R) in meters - Rmin=2.0, # minor radius (a) in meters Ai=2.5, # amu of main ion (if multiple isotope, make average) - B0=5.3, # Toroidal magnetic field on axis [T] Zeff=1.6, # needed for qlknn and fusion power # effective impurity charge state assumed for matching dilution=0.862. Zimp=10, @@ -167,6 +164,9 @@ def get_geometry(config: config_lib.Config) -> geometry.Geometry: config, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, # major radius (R) in meters + Rmin=2.0, # minor radius (a) in meters + B0=5.3, # Toroidal magnetic field on axis [T] ) diff --git a/torax/tests/test_data/test_iterhybrid_rampup.py b/torax/tests/test_data/test_iterhybrid_rampup.py index d619cfdb..a6e4afd9 100644 --- a/torax/tests/test_data/test_iterhybrid_rampup.py +++ b/torax/tests/test_data/test_iterhybrid_rampup.py @@ -39,10 +39,7 @@ def get_config() -> config_lib.Config: resistivity_mult=1, Ip={0: 3, 80: 10.5}, # total plasma current in MA # physical inputs - Rmaj=6.2, # major radius (R) in meters - Rmin=2.0, # minor radius (a) in meters Ai=2.5, # amu of main ion (if multiple isotope, make average) - B0=5.3, # Toroidal magnetic field on axis [T] Zeff=1.6, # needed for qlknn and fusion power # effective impurity charge state assumed for matching dilution=0.862. Zimp=10, @@ -186,6 +183,9 @@ def get_geometry(config: config_lib.Config) -> geometry.Geometry: config, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, # major radius (R) in meters + Rmin=2.0, # minor radius (a) in meters + B0=5.3, # Toroidal magnetic field on axis [T] ) diff --git a/torax/tests/test_lib/explicit_stepper.py b/torax/tests/test_lib/explicit_stepper.py index c6e38987..5bd910fc 100644 --- a/torax/tests/test_lib/explicit_stepper.py +++ b/torax/tests/test_lib/explicit_stepper.py @@ -125,7 +125,6 @@ def __call__( geo=geo, jtot_face=core_profiles_t.currents.jtot, psi=core_profiles_t.psi, - Rmaj=dynamic_config_slice_t.Rmaj, q_correction_factor=dynamic_config_slice_t.q_correction_factor, ) s_face = physics.calc_s_from_psi(geo, core_profiles_t.psi) diff --git a/torax/tests/test_lib/torax_refs.py b/torax/tests/test_lib/torax_refs.py index 736b8863..eb272e98 100644 --- a/torax/tests/test_lib/torax_refs.py +++ b/torax/tests/test_lib/torax_refs.py @@ -52,10 +52,7 @@ def circular_references() -> References: config, **{ 'nr': 25, - 'Rmaj': 6.2, - 'Rmin': 2.0, 'Ip': 15, - 'B0': 5.3, 'q_correction_factor': 1.0, 'nu': 3, 'fext': 0.2, @@ -67,6 +64,9 @@ def circular_references() -> References: config=config, kappa=1.72, hires_fac=4, + Rmaj=6.2, + Rmin=2.0, + B0=5.3, ) # ground truth values copied from example executions using # array.astype(str),which allows fully lossless reloading @@ -203,10 +203,7 @@ def chease_references_Ip_from_chease() -> References: # pylint: disable=invalid config, **{ 'nr': 25, - 'Rmaj': 6.2, - 'Rmin': 2.0, 'Ip': 15, - 'B0': 5.3, 'q_correction_factor': 1.0, 'nu': 3, 'fext': 0.2, @@ -219,6 +216,9 @@ def chease_references_Ip_from_chease() -> References: # pylint: disable=invalid geometry_dir=_GEO_DIRECTORY, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=False, + Rmaj=6.2, + Rmin=2.0, + B0=5.3, ) # ground truth values copied from an example PINT execution using # array.astype(str),which allows fully lossless reloading @@ -355,10 +355,7 @@ def chease_references_Ip_from_config() -> References: # pylint: disable=invalid config, **{ 'nr': 25, - 'Rmaj': 6.2, - 'Rmin': 2.0, 'Ip': 15, - 'B0': 5.3, 'q_correction_factor': 1.0, 'nu': 3, 'fext': 0.2, @@ -371,6 +368,9 @@ def chease_references_Ip_from_config() -> References: # pylint: disable=invalid geometry_dir=_GEO_DIRECTORY, geometry_file='ITER_hybrid_citrin_equil_cheasedata.mat2cols', Ip_from_parameters=True, + Rmaj=6.2, + Rmin=2.0, + B0=5.3, ) # ground truth values copied from an example executions using # array.astype(str),which allows fully lossless reloading diff --git a/torax/transport_model/critical_gradient.py b/torax/transport_model/critical_gradient.py index 52f89fa2..7b876c1c 100644 --- a/torax/transport_model/critical_gradient.py +++ b/torax/transport_model/critical_gradient.py @@ -82,11 +82,11 @@ def _call_implementation( (dynamic_config_slice.Ai * constants.mp) ** 0.5 / (constants.qe * geo.B0) ** 2 * (temp_ion_face * constants.keV2J) ** 1.5 - / dynamic_config_slice.Rmaj + / geo.Rmaj ) # R/LTi profile from current timestep temp_ion - rlti = -dynamic_config_slice.Rmaj * temp_ion_face_grad / temp_ion_face + rlti = -geo.Rmaj * temp_ion_face_grad / temp_ion_face # set minimum chi for PDE stability chi_ion = dynamic_config_slice.transport.chimin * jnp.ones_like( diff --git a/torax/transport_model/qlknn_wrapper.py b/torax/transport_model/qlknn_wrapper.py index a58ff1f8..dbce93a8 100644 --- a/torax/transport_model/qlknn_wrapper.py +++ b/torax/transport_model/qlknn_wrapper.py @@ -99,8 +99,6 @@ class _QLKNNRuntimeConfigInputs: """ # pylint: disable=invalid-name - Rmin: float - Rmaj: float nref: float Ai: float Zeff: float @@ -115,8 +113,6 @@ def from_config_slice( dynamic_config_slice: config_slice.DynamicConfigSlice, ) -> '_QLKNNRuntimeConfigInputs': return _QLKNNRuntimeConfigInputs( - Rmin=dynamic_config_slice.Rmin, - Rmaj=dynamic_config_slice.Rmaj, nref=dynamic_config_slice.nref, Ai=dynamic_config_slice.Ai, Zeff=dynamic_config_slice.Zeff, @@ -235,8 +231,8 @@ def _combined( model = _get_model() # pylint: disable=invalid-name - Rmin = runtime_config_inputs.Rmin - Rmaj = runtime_config_inputs.Rmaj + Rmin = geo.Rmin + Rmaj = geo.Rmaj # define radial coordinate as midplane average r # (typical assumption for transport models developed in circular geo) @@ -316,7 +312,6 @@ def _combined( geo, core_profiles.currents.jtot_face, core_profiles.psi, - Rmaj, runtime_config_inputs.q_correction_factor, ) smag = physics.calc_s_from_psi( @@ -345,7 +340,6 @@ def _combined( core_profiles=core_profiles, nref=runtime_config_inputs.nref, Zeff=runtime_config_inputs.Zeff, - Rmaj=Rmaj, coll_mult=runtime_config_inputs.transport.coll_mult, ) log_nu_star_face = jnp.log10(nu_star)