diff --git a/torax/calc_coeffs.py b/torax/calc_coeffs.py index c0385f64..0994eb5a 100644 --- a/torax/calc_coeffs.py +++ b/torax/calc_coeffs.py @@ -58,7 +58,6 @@ def calculate_pereverzev_flux( * true_ni_face * consts.keV2J * dynamic_runtime_params_slice.stepper.chi_per - / geo.rmax**2 ) chi_face_per_el = ( @@ -66,10 +65,9 @@ def calculate_pereverzev_flux( * true_ne_face * consts.keV2J * dynamic_runtime_params_slice.stepper.chi_per - / geo.rmax**2 ) - d_face_per_el = dynamic_runtime_params_slice.stepper.d_per / geo.rmax + d_face_per_el = dynamic_runtime_params_slice.stepper.d_per v_face_per_el = ( core_profiles.ne.face_grad() / core_profiles.ne.face_value() @@ -116,7 +114,7 @@ def calculate_pereverzev_flux( > dynamic_runtime_params_slice.profile_conditions.Ped_top, ), 0.0, - d_face_per_el * geo.g1_over_vpr_face / geo.rmax, + d_face_per_el * geo.g1_over_vpr_face, ) v_face_per_el = jnp.where( @@ -126,7 +124,7 @@ def calculate_pereverzev_flux( > dynamic_runtime_params_slice.profile_conditions.Ped_top, ), 0.0, - v_face_per_el * geo.g0_face / geo.rmax, + v_face_per_el * geo.g0_face, ) chi_face_per_ion = chi_face_per_ion.at[0].set(chi_face_per_ion[1]) @@ -404,9 +402,7 @@ def _calc_coeffs_full( chi_face_el = transport_coeffs.chi_face_el d_face_el = transport_coeffs.d_face_el v_face_el = transport_coeffs.v_face_el - # TODO(b/351356977): remove rmax from TORAX by normalizing other quantities - # like g2g3_over_rho_face, vpr, spr - d_face_psi = geo.g2g3_over_rho_face * geo.rmax + d_face_psi = geo.g2g3_over_rhon_face if static_runtime_params_slice.dens_eq: if d_face_el is None or v_face_el is None: @@ -533,19 +529,17 @@ def _calc_coeffs_full( * true_ni_face * consts.keV2J * chi_face_ion - / geo.rmax**2 ) full_chi_face_el = ( geo.g1_over_vpr_face * true_ne_face * consts.keV2J * chi_face_el - / geo.rmax**2 ) # entire coefficient preceding dne/dr in particle equation - full_d_face_el = geo.g1_over_vpr_face * d_face_el / geo.rmax**2 - full_v_face_el = geo.g0_face * v_face_el / geo.rmax + full_d_face_el = geo.g1_over_vpr_face * d_face_el + full_v_face_el = geo.g0_face * v_face_el # density source terms. Initialize source matrix to zero source_mat_nn = jnp.zeros_like(geo.r) @@ -624,7 +618,6 @@ def _calc_coeffs_full( * geo.vpr_face * true_ni_face * consts.keV2J - / geo.rmax ) v_heat_face_el += ( @@ -636,7 +629,6 @@ def _calc_coeffs_full( * geo.vpr_face * true_ne_face * consts.keV2J - / geo.rmax ) # Add phibdot terms to particle transport convection @@ -647,7 +639,6 @@ def _calc_coeffs_full( / geo.Phib * geo.r_face_norm * geo.vpr_face - / geo.rmax ) # Add phibdot terms to poloidal flux convection @@ -660,7 +651,6 @@ def _calc_coeffs_full( * sigma_face * geo.r_face_norm**2 / geo.F_face**2 - / geo.rmax ) # Ion and electron heat sources. @@ -748,7 +738,7 @@ def _calc_coeffs_full( # Add effective phibdot heat source terms # second derivative of volume profile with respect to r_norm - vprpr_norm = math_utils.gradient(geo.vpr, geo.r) / geo.rmax**2 + vprpr_norm = math_utils.gradient(geo.vpr, geo.r_norm) source_i += ( 1.0 diff --git a/torax/config/runtime_params.py b/torax/config/runtime_params.py index 77085367..9bb75d4a 100644 --- a/torax/config/runtime_params.py +++ b/torax/config/runtime_params.py @@ -185,10 +185,10 @@ class Numerics: # numerical (e.g. no. of grid points, other info needed by solver) # effective source to dominate PDE in internal boundary condtion location # if T != Tped - largeValue_T: float = 1.0e10 + largeValue_T: float = 2.0e10 # effective source to dominate density PDE in internal boundary condtion # location if n != neped - largeValue_n: float = 1.0e8 + largeValue_n: float = 2.0e8 # NOMUTANTS -- It's expected for the tests to pass with different defaults. diff --git a/torax/core_profile_setters.py b/torax/core_profile_setters.py index 38bd25fe..4d355229 100644 --- a/torax/core_profile_setters.py +++ b/torax/core_profile_setters.py @@ -235,7 +235,7 @@ def _prescribe_currents_no_bootstrap( 1 - geo.r_face_norm**2 ) ** dynamic_runtime_params_slice.profile_conditions.nu # calculate total and Ohmic current profiles - denom = _trapz(jformula_face * geo.spr_face, geo.r_face) + denom = _trapz(jformula_face * geo.spr_face, geo.r_face_norm) if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current: Ctot = Ip * 1e6 / denom jtot_face = jformula_face * Ctot @@ -345,7 +345,7 @@ def _prescribe_currents_with_bootstrap( jformula_face = ( 1 - geo.r_face_norm**2 ) ** dynamic_runtime_params_slice.profile_conditions.nu - denom = _trapz(jformula_face * geo.spr_face, geo.r_face) + denom = _trapz(jformula_face * geo.spr_face, geo.r_face_norm) # calculate total and Ohmic current profiles if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current: Ctot = Ip * 1e6 / denom @@ -495,8 +495,8 @@ def _update_psi_from_j( psi_constraint = ( dynamic_runtime_params_slice.profile_conditions.Ip * 1e6 - * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0) - / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1]) + * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax) + / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1]) * geo.rmax ) @@ -504,12 +504,12 @@ def _update_psi_from_j( assert y.ndim == 1 assert geo.r_hires.ndim == 1 integrated = math_utils.cumulative_trapezoid( - geo.r_hires, y, initial=jnp.zeros(()) + geo.r_hires_norm, y, initial=jnp.zeros(()) ) scale = jnp.concatenate(( jnp.zeros((1,)), - (8 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.B0) - / (geo.F_hires[1:] * geo.Rmaj * geo.g2g3_over_rho_hires[1:]), + (8 * jnp.pi**3 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax) + / (geo.F_hires[1:] * geo.Rmaj * geo.g2g3_over_rhon_hires[1:]), )) # dpsi_dr on the cell grid dpsi_dr_hires = scale * integrated @@ -612,8 +612,8 @@ def initial_core_profiles( psi_constraint = ( dynamic_runtime_params_slice.profile_conditions.Ip * 1e6 - * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0) - / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1]) + * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax) + / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1]) * geo.rmax ) psi = cell_variable.CellVariable( @@ -873,8 +873,8 @@ def compute_boundary_conditions( 'psi': dict( right_face_grad_constraint=Ip * 1e6 - * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0) - / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1]) + * (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax) + / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1]) * geo.rmax, right_face_constraint=None, ), @@ -906,7 +906,7 @@ def _get_jtot_hires( jformula_hires = ( 1 - geo.r_hires_norm**2 ) ** dynamic_runtime_params_slice.profile_conditions.nu - denom = _trapz(jformula_hires * geo.spr_hires, geo.r_hires) + denom = _trapz(jformula_hires * geo.spr_hires, geo.r_hires_norm) if dynamic_runtime_params_slice.profile_conditions.initial_j_is_total_current: Ctot_hires = ( dynamic_runtime_params_slice.profile_conditions.Ip * 1e6 / denom diff --git a/torax/examples/iterhybrid_predictor_corrector.py b/torax/examples/iterhybrid_predictor_corrector.py index 3a3c3338..5331feef 100644 --- a/torax/examples/iterhybrid_predictor_corrector.py +++ b/torax/examples/iterhybrid_predictor_corrector.py @@ -68,12 +68,6 @@ # likely be increased further beyond this default. 'dtmult': 50, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary\ - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/examples/iterhybrid_rampup.py b/torax/examples/iterhybrid_rampup.py index 5ee04b83..a639187d 100644 --- a/torax/examples/iterhybrid_rampup.py +++ b/torax/examples/iterhybrid_rampup.py @@ -77,12 +77,6 @@ # likely be increased further beyond this default. 'dtmult': 30, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/geometry.py b/torax/geometry.py index d81e57de..9e70a397 100644 --- a/torax/geometry.py +++ b/torax/geometry.py @@ -153,9 +153,9 @@ class Geometry: g2_face: chex.Array g3: chex.Array g3_face: chex.Array - g2g3_over_rho: chex.Array - g2g3_over_rho_face: chex.Array - g2g3_over_rho_hires: chex.Array + g2g3_over_rhon: chex.Array + g2g3_over_rhon_face: chex.Array + g2g3_over_rhon_hires: chex.Array F: chex.Array F_face: chex.Array F_hires: chex.Array @@ -207,7 +207,7 @@ def g1_over_vpr2(self) -> chex.Array: @property def g0_over_vpr_face(self) -> jax.Array: return jnp.concatenate(( - jnp.ones(1), # correct value is unity on-axis + jnp.ones(1) / self.rmax, # correct value is 1/rmax on-axis self.g0_face[1:] / self.vpr_face[1:], # avoid div by zero on-axis )) @@ -221,7 +221,7 @@ def g1_over_vpr_face(self) -> jax.Array: @property def g1_over_vpr2_face(self) -> jax.Array: return jnp.concatenate(( - jnp.ones(1), # correct value is unity on-axis + jnp.ones(1) / self.rmax**2, # correct value is 1/rmax**2 on-axis self.g1_face[1:] / self.vpr_face[1:] ** 2, # avoid div by zero on-axis )) @@ -254,9 +254,9 @@ class GeometryProvider: g2_face: interpolated_param.InterpolatedVarSingleAxis g3: interpolated_param.InterpolatedVarSingleAxis g3_face: interpolated_param.InterpolatedVarSingleAxis - g2g3_over_rho: interpolated_param.InterpolatedVarSingleAxis - g2g3_over_rho_face: interpolated_param.InterpolatedVarSingleAxis - g2g3_over_rho_hires: interpolated_param.InterpolatedVarSingleAxis + g2g3_over_rhon: interpolated_param.InterpolatedVarSingleAxis + g2g3_over_rhon_face: interpolated_param.InterpolatedVarSingleAxis + g2g3_over_rhon_hires: interpolated_param.InterpolatedVarSingleAxis F: interpolated_param.InterpolatedVarSingleAxis F_face: interpolated_param.InterpolatedVarSingleAxis F_hires: interpolated_param.InterpolatedVarSingleAxis @@ -455,21 +455,19 @@ def build_circular_geometry( area = np.pi * r**2 * kappa area_face = np.pi * r_face**2 * kappa_face - # V' for volume integrations - vpr = ( - 4 * np.pi**2 * Rmaj * r * kappa - + volume / kappa * (kappa_param - 1) / rmax + # V' = dV/drnorm for volume integrations + vpr = 4 * np.pi**2 * Rmaj * r * kappa * rmax + volume / kappa * ( + kappa_param - 1 ) vpr_face = ( - 4 * np.pi**2 * Rmaj * r_face * kappa_face - + volume_face / kappa_face * (kappa_param - 1) / rmax + 4 * np.pi**2 * Rmaj * r_face * kappa_face * rmax + + volume_face / kappa_face * (kappa_param - 1) ) # pylint: disable=invalid-name - # S' for area integrals on cell grid - spr_cell = 2 * np.pi * r * kappa + area / kappa * (kappa_param - 1) / rmax - spr_face = ( - 2 * np.pi * r_face * kappa_face - + area_face / kappa_face * (kappa_param - 1) / rmax + # S' = dS/drnorm for area integrals on cell grid + spr_cell = 2 * np.pi * r * kappa * rmax + area / kappa * (kappa_param - 1) + spr_face = 2 * np.pi * r_face * kappa_face * rmax + area_face / kappa_face * ( + kappa_param - 1 ) delta_face = np.zeros(len(r_face)) @@ -478,12 +476,12 @@ def build_circular_geometry( # With circular geometry approximation. # g0: <\nabla V> - g0 = vpr - g0_face = vpr_face + g0 = vpr / rmax + g0_face = vpr_face / rmax - # g1: <\nabla V^2> - g1 = vpr**2 - g1_face = vpr_face**2 + # g1: <(\nabla V)^2> + g1 = vpr**2 / rmax**2 + g1_face = vpr_face**2 / rmax**2 # g2: <(\nabla V)^2 / R^2> # V = 2*pi^2*R*r^2*kappa @@ -519,15 +517,15 @@ def build_circular_geometry( F_face = np.ones(len(r_face)) * Rmaj * B0 # Using an approximation where: - # g2g3_over_rho = 16 * pi**4 * G2 / (J * R) where: + # g2g3_over_rhon = 16 * pi**4 * G2 / (J * R) where: # G2 = vpr / (4 * pi**2) * <1/R^2> # This is done due to our ad-hoc kappa assumption, which leads to more - # reasonable values for g2g3_over_rho through the G2 definition. + # reasonable values for g2g3_over_rhon through the G2 definition. # In the future, a more rigorous analytical geometry will be developed and - # the direct definition of g2g3_over_rho will be used. + # the direct definition of g2g3_over_rhon will be used. - g2g3_over_rho = 4 * np.pi**2 * vpr * g3 / (J * Rmaj) - g2g3_over_rho_face = 4 * np.pi**2 * vpr_face * g3_face / (J_face * Rmaj) + g2g3_over_rhon = 4 * np.pi**2 * vpr * g3 / (J * Rmaj) + g2g3_over_rhon_face = 4 * np.pi**2 * vpr_face * g3_face / (J_face * Rmaj) # High resolution versions for j (plasma current) and psi (poloidal flux) # manipulations. Needed if psi is initialized from plasma current, which is @@ -547,20 +545,20 @@ def build_circular_geometry( volume_hires = 2 * np.pi**2 * Rmaj * r_hires**2 * kappa_hires area_hires = np.pi * r_hires**2 * kappa_hires - # V' for volume integrations on hires grid + # V' = dV/drnorm for volume integrations on hires grid vpr_hires = ( - 4 * np.pi**2 * Rmaj * r_hires * kappa_hires - + volume_hires / kappa_hires * (kappa_param - 1) / rmax + 4 * np.pi**2 * Rmaj * r_hires * kappa_hires * rmax + + volume_hires / kappa_hires * (kappa_param - 1) ) - # S' for area integrals on hires grid + # S' = dS/drnorm for area integrals on hires grid spr_hires = ( - 2 * np.pi * r_hires * kappa_hires - + area_hires / kappa_hires * (kappa_param - 1) / rmax + 2 * np.pi * r_hires * kappa_hires * rmax + + area_hires / kappa_hires * (kappa_param - 1) ) g3_hires = 1 / (Rmaj**2 * (1 - (r_hires / Rmaj) ** 2) ** (3.0 / 2.0)) F_hires = np.ones(len(r_hires)) * B0 * Rmaj - g2g3_over_rho_hires = 4 * np.pi**2 * vpr_hires * g3_hires * B0 / F_hires + g2g3_over_rhon_hires = 4 * np.pi**2 * vpr_hires * g3_hires * B0 / F_hires return CircularAnalyticalGeometry( # Set the standard geometry params. @@ -588,9 +586,9 @@ def build_circular_geometry( g2_face=g2_face, g3=g3, g3_face=g3_face, - g2g3_over_rho=g2g3_over_rho, - g2g3_over_rho_face=g2g3_over_rho_face, - g2g3_over_rho_hires=g2g3_over_rho_hires, + g2g3_over_rhon=g2g3_over_rhon, + g2g3_over_rhon_face=g2g3_over_rhon_face, + g2g3_over_rhon_hires=g2g3_over_rhon_hires, F=F, F_face=F_face, F_hires=F_hires, @@ -644,7 +642,7 @@ class StandardGeometryIntermediates: rhon: Normalized toroidal flux coordinate Rin: Radius of the flux surface at the inboard side at midplane Rout: Radius of the flux surface at the outboard side at midplane - RBphi: Toroidal field flux function + F: Toroidal field flux function int_dl_over_Bp: oint (dl / Bp) (contour integral) flux_surf_avg_1_over_R2: <1/R**2> flux_surf_avg_Bp2: @@ -665,11 +663,10 @@ class StandardGeometryIntermediates: B: chex.Numeric psi: chex.Array Ip_profile: chex.Array - rho: chex.Array - rhon: np.ndarray + Phi: chex.Array Rin: chex.Array Rout: chex.Array - RBphi: chex.Array + F: chex.Array int_dl_over_Bp: chex.Array flux_surf_avg_1_over_R2: chex.Array flux_surf_avg_Bp2: chex.Array @@ -729,14 +726,14 @@ def from_chease( psi = chease_data['PSIchease=psi/2pi'] * psiunnormfactor * 2 * np.pi Ip_chease = chease_data['Ipprofile'] / constants.CONSTANTS.mu0 * Rmaj * B0 - # toroidal flux coordinate - rho = chease_data['RHO_TOR=sqrt(Phi/pi/B0)'] * Rmaj - rhon = chease_data['RHO_TOR_NORM'] + # toroidal flux + Phi = (chease_data['RHO_TOR=sqrt(Phi/pi/B0)'] * Rmaj) ** 2 * B0 * np.pi + # midplane radii Rin_chease = chease_data['R_INBOARD'] * Rmaj Rout_chease = chease_data['R_OUTBOARD'] * Rmaj # toroidal field flux function - RBphi = chease_data['T=RBphi'] * Rmaj * B0 + F = chease_data['T=RBphi'] * Rmaj * B0 int_dl_over_Bp = chease_data['Int(Rdlp/|grad(psi)|)=Int(Jdchi)'] * Rmaj / B0 flux_surf_avg_1_over_R2 = chease_data['<1/R**2>'] / Rmaj**2 @@ -757,11 +754,10 @@ def from_chease( B=B0, psi=psi, Ip_profile=Ip_chease, - rho=rho, - rhon=rhon, + Phi=Phi, Rin=Rin_chease, Rout=Rout_chease, - RBphi=RBphi, + F=F, int_dl_over_Bp=int_dl_over_Bp, flux_surf_avg_1_over_R2=flux_surf_avg_1_over_R2, flux_surf_avg_Bp2=flux_surf_avg_Bp2, @@ -789,6 +785,11 @@ def build_standard_geometry( Returns: A StandardGeometry object. """ + + # Toroidal flux coordinates + rho = np.sqrt(intermediate.Phi / (np.pi * intermediate.B)) + rhon = rho / rho[-1] + # flux surface integrals of various geometry quantities C1 = intermediate.int_dl_over_Bp @@ -804,36 +805,34 @@ def build_standard_geometry( g2 = C1 * C3 * 4 * np.pi**2 # <(\nabla psi)**2 / R**2> * (dV/dpsi) ** 2 g3 = C2[1:] / C1[1:] # <1/R**2> g3 = np.concatenate((np.array([1 / intermediate.Rin[0] ** 2]), g3)) - g2g3_over_rho = g2[1:] * g3[1:] / intermediate.rho[1:] - g2g3_over_rho = np.concatenate((np.zeros(1), g2g3_over_rho)) - - J = intermediate.RBphi / (intermediate.Rmaj * intermediate.B) + g2g3_over_rhon = g2[1:] * g3[1:] / rhon[1:] + g2g3_over_rhon = np.concatenate((np.zeros(1), g2g3_over_rhon)) # make an alternative initial psi, self-consistent with CHEASE Ip profile # needed because CHEASE psi profile has noisy second derivatives dpsidrho = ( intermediate.Ip_profile[1:] - * (16 * constants.CONSTANTS.mu0 * np.pi**4) - / (g2g3_over_rho[1:] * intermediate.Rmaj * J[1:]) + * (16 * constants.CONSTANTS.mu0 * np.pi**4 * rho[-1] * intermediate.B) + / (g2g3_over_rhon[1:] * intermediate.F[1:]) ) dpsidrho = np.concatenate((np.zeros(1), dpsidrho)) psi_from_Ip = scipy.integrate.cumulative_trapezoid( - y=dpsidrho, x=intermediate.rho, initial=0.0 + y=dpsidrho, x=rho, initial=0.0 ) # set Ip-consistent psi derivative boundary condition (although will be # replaced later with an fvm constraint) psi_from_Ip[-1] = psi_from_Ip[-2] + ( - 16 * constants.CONSTANTS.mu0 * np.pi**4 + 16 * constants.CONSTANTS.mu0 * np.pi**4 * rho[-1] * intermediate.B ) * intermediate.Ip_profile[-1] / ( - g2g3_over_rho[-1] * intermediate.Rmaj * J[-1] + g2g3_over_rhon[-1] * intermediate.F[-1] ) * ( - intermediate.rho[-1] - intermediate.rho[-2] + rho[-1] - rho[-2] ) - # dV/drho, dS/drho - vpr = np.gradient(intermediate.volume, intermediate.rho) - spr = np.gradient(intermediate.area, intermediate.rho) + # dV/drhon, dS/drhon + vpr = np.gradient(intermediate.volume, rhon) + spr = np.gradient(intermediate.area, rhon) # gradient boundary approximation not appropriate here vpr[0] = 0 spr[0] = 0 @@ -851,10 +850,10 @@ def build_standard_geometry( # fill geometry structure # r_norm coordinate is rho_tor_norm - dr_norm = float(intermediate.rhon[-1]) / intermediate.nr + dr_norm = float(rhon[-1]) / intermediate.nr # normalized grid mesh = Grid1D.construct(nx=intermediate.nr, dx=dr_norm) - rmax = intermediate.rho[-1] # radius denormalization constant + rmax = rho[-1] # radius denormalization constant # helper variables for mesh cells and faces r_face_norm = mesh.face_centers r_norm = mesh.cell_centers @@ -864,7 +863,7 @@ def build_standard_geometry( r_hires_norm = np.linspace(0, 1, intermediate.nr * intermediate.hires_fac) r_hires = r_hires_norm * rmax - rhon_interpolation_func = lambda x, y: np.interp(x, intermediate.rhon, y) + rhon_interpolation_func = lambda x, y: np.interp(x, rhon, y) # V' for volume integrations on face grid vpr_face = rhon_interpolation_func(r_face_norm, vpr) # V' for volume integrations on cell grid @@ -888,9 +887,9 @@ def build_standard_geometry( # average triangularity delta_face = 0.5 * (delta_upper_face + delta_lower_face) - F_face = rhon_interpolation_func(r_face_norm, intermediate.RBphi) - F = rhon_interpolation_func(r_norm, intermediate.RBphi) - F_hires = rhon_interpolation_func(r_hires_norm, intermediate.RBphi) + F_face = rhon_interpolation_func(r_face_norm, intermediate.F) + F = rhon_interpolation_func(r_norm, intermediate.F) + F_hires = rhon_interpolation_func(r_hires_norm, intermediate.F) psi = rhon_interpolation_func(r_norm, intermediate.psi) psi_from_Ip = rhon_interpolation_func(r_norm, psi_from_Ip) @@ -916,9 +915,9 @@ def build_standard_geometry( g3_face = rhon_interpolation_func(r_face_norm, g3) g3 = rhon_interpolation_func(r_norm, g3) - g2g3_over_rho_face = rhon_interpolation_func(r_face_norm, g2g3_over_rho) - g2g3_over_rho_hires = rhon_interpolation_func(r_hires_norm, g2g3_over_rho) - g2g3_over_rho = rhon_interpolation_func(r_norm, g2g3_over_rho) + g2g3_over_rhon_face = rhon_interpolation_func(r_face_norm, g2g3_over_rhon) + g2g3_over_rhon_hires = rhon_interpolation_func(r_hires_norm, g2g3_over_rhon) + g2g3_over_rhon = rhon_interpolation_func(r_norm, g2g3_over_rhon) volume_face = rhon_interpolation_func(r_face_norm, intermediate.volume) volume_hires = rhon_interpolation_func(r_hires_norm, intermediate.volume) @@ -953,9 +952,9 @@ def build_standard_geometry( g2_face=g2_face, g3=g3, g3_face=g3_face, - g2g3_over_rho=g2g3_over_rho, - g2g3_over_rho_face=g2g3_over_rho_face, - g2g3_over_rho_hires=g2g3_over_rho_hires, + g2g3_over_rhon=g2g3_over_rhon, + g2g3_over_rhon_face=g2g3_over_rhon_face, + g2g3_over_rhon_hires=g2g3_over_rhon_hires, F=F, F_face=F_face, F_hires=F_hires, diff --git a/torax/physics.py b/torax/physics.py index 9b7d0900..fcd44d1e 100644 --- a/torax/physics.py +++ b/torax/physics.py @@ -201,10 +201,10 @@ def calc_jtot_from_psi( # pylint: disable=invalid-name I_tot = ( dpsi_dr - * geo.g2g3_over_rho_face + * geo.g2g3_over_rhon_face * geo.F_face / geo.B0 - / (16 * jnp.pi**4 * constants.CONSTANTS.mu0) + / (16 * jnp.pi**4 * constants.CONSTANTS.mu0 * geo.rmax) ) jtot_face = ( @@ -238,12 +238,12 @@ def calc_s_from_psi( # Volume on face grid # pylint:disable=invalid-name - V = math_utils.cumulative_trapezoid(geo.r_face, geo.vpr_face) + V = math_utils.cumulative_trapezoid(geo.r_face, geo.vpr_face * geo.rmax) s = ( 2 * V - / (geo.r_face[1:] * geo.vpr_face[1:]) + / (geo.r_face[1:] * geo.vpr_face[1:] * geo.rmax) * (1 - geo.r_face[1:] * d2psi_d2r / dpsi_dr[1:]) ) s = jnp.concatenate((s[0:1], s)) diff --git a/torax/sources/bootstrap_current_source.py b/torax/sources/bootstrap_current_source.py index f50c9a5b..81f81c05 100644 --- a/torax/sources/bootstrap_current_source.py +++ b/torax/sources/bootstrap_current_source.py @@ -395,7 +395,7 @@ def calc_neoclassical( I_bootstrap = integrate.trapezoid( j_bootstrap_face * geo.spr_face, - geo.r_face, + geo.r_face_norm, ) return source_profiles.BootstrapCurrentProfile( diff --git a/torax/sources/bremsstrahlung_heat_sink.py b/torax/sources/bremsstrahlung_heat_sink.py index 1afe1837..ba0b6f6a 100644 --- a/torax/sources/bremsstrahlung_heat_sink.py +++ b/torax/sources/bremsstrahlung_heat_sink.py @@ -104,7 +104,7 @@ def calc_relativistic_correction() -> jax.Array: # In MW P_brem_total = jax.scipy.integrate.trapezoid( - P_brem_profile_face * geo.vpr_face, geo.r_face + P_brem_profile_face * geo.vpr_face, geo.r_face_norm ) return P_brem_total, P_brem_profile_cell diff --git a/torax/sources/external_current_source.py b/torax/sources/external_current_source.py index 0a40b9fb..f961e758 100644 --- a/torax/sources/external_current_source.py +++ b/torax/sources/external_current_source.py @@ -120,7 +120,7 @@ def _calculate_jext_face( # pytype: disable=name-error / (2 * dynamic_source_runtime_params.wext**2) ) - Cext = Iext * 1e6 / _trapz(jextform_face * geo.spr_face, geo.r_face) + Cext = Iext * 1e6 / _trapz(jextform_face * geo.spr_face, geo.r_face_norm) jext_face = Cext * jextform_face # external current profile return jext_face @@ -158,7 +158,9 @@ def _calculate_jext_hires( # pytype: disable=name-error -((geo.r_hires_norm - dynamic_source_runtime_params.rext) ** 2) / (2 * dynamic_source_runtime_params.wext**2) ) - Cext_hires = Iext * 1e6 / _trapz(jextform_hires * geo.spr_hires, geo.r_hires) + Cext_hires = ( + Iext * 1e6 / _trapz(jextform_hires * geo.spr_hires, geo.r_hires_norm) + ) # External current profile on cell grid jext_hires = Cext_hires * jextform_hires return jext_hires diff --git a/torax/sources/formulas.py b/torax/sources/formulas.py index 69a7cedd..31e87423 100644 --- a/torax/sources/formulas.py +++ b/torax/sources/formulas.py @@ -44,10 +44,10 @@ def exponential_profile( | cell = exp(-(c1 - r) / c2) | face = exp(-(c1 - r_face) / c2) - | C = total / trapz(vpr_face * face, r_face) + | C = total / trapz(vpr_face * face, r_face_norm) | profile = C * cell - The formula can use the normalized r and r_face if specified. + The formula can use the normalized r and r_face for c1 + c2 if specified. Args: geo: Geometry constants of torus. @@ -65,7 +65,9 @@ def exponential_profile( S = jnp.exp(-(c1 - r) / c2) S_face = jnp.exp(-(c1 - r_face) / c2) # calculate constant prefactor - C = total / jax.scipy.integrate.trapezoid(geo.vpr_face * S_face, geo.r_face) + C = total / jax.scipy.integrate.trapezoid( + geo.vpr_face * S_face, geo.r_face_norm + ) return C * S @@ -82,10 +84,10 @@ def gaussian_profile( | cell = exp(-( (r - c1)**2 / (2 * c2**2) )) | face = exp(-( (r_face - c1)**2 / (2 * c2**2) )) - | C = total / trazp(vpr_face * face, r_face) + | C = total / trapz( vpr_face * face, r_face_norm) | profile = C * cell - The formula can use the normalized r and r_face if specified. + The formula can use the normalized r and r_face for c1 + c2 if specified. Args: geo: Geometry constants of torus. @@ -103,7 +105,9 @@ def gaussian_profile( S = jnp.exp(-((r - c1) ** 2) / (2 * c2**2)) S_face = jnp.exp(-((r_face - c1) ** 2) / (2 * c2**2)) # calculate constant prefactor - C = total / jax.scipy.integrate.trapezoid(geo.vpr_face * S_face, geo.r_face) + C = total / jax.scipy.integrate.trapezoid( + geo.vpr_face * S_face, geo.r_face_norm + ) return C * S diff --git a/torax/sources/fusion_heat_source.py b/torax/sources/fusion_heat_source.py index 762c1c32..96c65c22 100644 --- a/torax/sources/fusion_heat_source.py +++ b/torax/sources/fusion_heat_source.py @@ -96,7 +96,8 @@ def calc_fusion( # [MW] Ptot = ( - jax.scipy.integrate.trapezoid(Pfus_face * geo.vpr_face, geo.r_face) / 1e6 + jax.scipy.integrate.trapezoid(Pfus_face * geo.vpr_face, geo.r_face_norm) + / 1e6 ) alpha_fraction = 3.5 / 17.6 # fusion power fraction to alpha particles diff --git a/torax/sources/generic_ion_el_heat_source.py b/torax/sources/generic_ion_el_heat_source.py index d82fbc58..50f27548 100644 --- a/torax/sources/generic_ion_el_heat_source.py +++ b/torax/sources/generic_ion_el_heat_source.py @@ -94,7 +94,9 @@ def calc_generic_heat_source( Q = jnp.exp(-((geo.r_norm - rsource) ** 2) / (2 * w**2)) Q_face = jnp.exp(-((geo.r_face_norm - rsource) ** 2) / (2 * w**2)) # calculate constant prefactor - C = Ptot / jax.scipy.integrate.trapezoid(geo.vpr_face * Q_face, geo.r_face) + C = Ptot / jax.scipy.integrate.trapezoid( + geo.vpr_face * Q_face, geo.r_face_norm + ) source_ion = C * Q * (1 - el_heat_fraction) source_el = C * Q * el_heat_fraction diff --git a/torax/sources/source_models.py b/torax/sources/source_models.py index a13e9017..64657fa5 100644 --- a/torax/sources/source_models.py +++ b/torax/sources/source_models.py @@ -378,7 +378,7 @@ def sum_sources_psi( ) mu0 = constants.CONSTANTS.mu0 prefactor = ( - 8 * geo.vpr * geo.rmax * jnp.pi**2 * geo.B0 * mu0 * geo.Phib / geo.F**2 + 8 * geo.vpr * jnp.pi**2 * geo.B0 * mu0 * geo.Phib / geo.F**2 ) scale_source = lambda src: -src * prefactor return scale_source(total) @@ -468,7 +468,7 @@ def calc_and_sum_sources_psi( mu0 = constants.CONSTANTS.mu0 prefactor = ( - 8 * geo.vpr * geo.rmax * jnp.pi**2 * geo.B0 * mu0 * geo.Phib / geo.F**2 + 8 * geo.vpr * jnp.pi**2 * geo.B0 * mu0 * geo.Phib / geo.F**2 ) scale_source = lambda src: -src * prefactor @@ -525,7 +525,7 @@ def calc_psidot( / geo.F**2 ) d_face_psi = ( - geo.g2g3_over_rho_face * geo.rmax + geo.g2g3_over_rhon_face ) c_mat, c = diffusion_terms.make_diffusion_terms(d_face_psi, core_profiles.psi) diff --git a/torax/sources/tests/fusion_heat_source.py b/torax/sources/tests/fusion_heat_source.py index c228f2d6..d2759872 100644 --- a/torax/sources/tests/fusion_heat_source.py +++ b/torax/sources/tests/fusion_heat_source.py @@ -121,7 +121,7 @@ def calculate_fusion(runtime_params, geo, core_profiles): * (core_profiles.ni.face_value() * runtime_params.numerics.nref) ** 2 * sigmav ) # [W/m^3] - Ptot = np.trapz(Pfus * geo.vpr_face, geo.r_face) / 1e6 # [MW] + Ptot = np.trapz(Pfus * geo.vpr_face, geo.r_face_norm) / 1e6 # [MW] return Ptot diff --git a/torax/tests/boundary_conditions.py b/torax/tests/boundary_conditions.py index 5f033486..f288ae33 100644 --- a/torax/tests/boundary_conditions.py +++ b/torax/tests/boundary_conditions.py @@ -100,8 +100,8 @@ def test_setting_boundary_conditions( psi_constraint = ( 6e6 - * (16 * np.pi**4 * constants.CONSTANTS.mu0 * geo.B0) - / (geo.g2g3_over_rho_face[-1] * geo.F_face[-1]) + * (16 * np.pi**4 * constants.CONSTANTS.mu0 * geo.B0 * geo.rmax) + / (geo.g2g3_over_rhon_face[-1] * geo.F_face[-1]) * geo.rmax ) np.testing.assert_allclose(updated.temp_ion.right_face_constraint, 27.7) diff --git a/torax/tests/geometry.py b/torax/tests/geometry.py index ce92bf34..5563265d 100644 --- a/torax/tests/geometry.py +++ b/torax/tests/geometry.py @@ -82,11 +82,10 @@ def foo(geo: geometry.Geometry): # Use the same dummy value for the rest. psi=np.arange(0, 1.0, 0.01), Ip_profile=np.arange(0, 1.0, 0.01), - rho=np.arange(0, 1.0, 0.01), - rhon=np.arange(0, 1.0, 0.01), + Phi=np.arange(0, 1.0, 0.01), Rin=np.arange(0, 1.0, 0.01), Rout=np.arange(0, 1.0, 0.01), - RBphi=np.arange(0, 1.0, 0.01), + F=np.arange(0, 1.0, 0.01), int_dl_over_Bp=np.arange(0, 1.0, 0.01), flux_surf_avg_1_over_R2=np.arange(0, 1.0, 0.01), flux_surf_avg_Bp2=np.arange(0, 1.0, 0.01), @@ -117,11 +116,10 @@ def test_build_geometry_provider(self): # Use the same dummy value for the rest. psi=np.arange(0, 1.0, 0.01), Ip_profile=np.arange(0, 1.0, 0.01), - rho=np.arange(0, 1.0, 0.01), - rhon=np.arange(0, 1.0, 0.01), + Phi=np.arange(0, 1.0, 0.01), Rin=np.arange(0, 1.0, 0.01), Rout=np.arange(0, 1.0, 0.01), - RBphi=np.arange(0, 1.0, 0.01), + F=np.arange(0, 1.0, 0.01), int_dl_over_Bp=np.arange(0, 1.0, 0.01), flux_surf_avg_1_over_R2=np.arange(0, 1.0, 0.01), flux_surf_avg_Bp2=np.arange(0, 1.0, 0.01), @@ -144,11 +142,10 @@ def test_build_geometry_provider(self): # Use the same dummy value for the rest. psi=np.arange(0, 1.0, 0.01), Ip_profile=np.arange(0, 2.0, 0.02), - rho=np.arange(0, 1.0, 0.01), - rhon=np.arange(0, 1.0, 0.01), + Phi=np.arange(0, 1.0, 0.01), Rin=np.arange(0, 1.0, 0.01), Rout=np.arange(0, 1.0, 0.01), - RBphi=np.arange(0, 1.0, 0.01), + F=np.arange(0, 1.0, 0.01), int_dl_over_Bp=np.arange(0, 1.0, 0.01), flux_surf_avg_1_over_R2=np.arange(0, 1.0, 0.01), flux_surf_avg_Bp2=np.arange(0, 1.0, 0.01), diff --git a/torax/tests/sim.py b/torax/tests/sim.py index 65f5a742..2adb7cf2 100644 --- a/torax/tests/sim.py +++ b/torax/tests/sim.py @@ -333,7 +333,7 @@ class SimTest(sim_test_case.SimTestCase): 'test_iterhybrid_rampup.py', _ALL_PROFILES, 0, - 2e-7 + 1e-6 ), # Tests time-dependent circular geometry. ( diff --git a/torax/tests/state.py b/torax/tests/state.py index e28c29bf..207cc2d0 100644 --- a/torax/tests/state.py +++ b/torax/tests/state.py @@ -308,7 +308,7 @@ def test_initial_psi_from_j( # calculate total and Ohmic current profiles arising from nu=2 jformula_face = (1 - geo.r_face_norm**2) ** 2 denom = jax.scipy.integrate.trapezoid( - jformula_face * geo.spr_face, geo.r_face + jformula_face * geo.spr_face, geo.r_face_norm ) ctot = config1.profile_conditions.Ip * 1e6 / denom jtot_formula_face = jformula_face * ctot diff --git a/torax/tests/test_data/test_all_transport_crank_nicolson.nc b/torax/tests/test_data/test_all_transport_crank_nicolson.nc index 0ceb7309..612bf327 100644 Binary files a/torax/tests/test_data/test_all_transport_crank_nicolson.nc and b/torax/tests/test_data/test_all_transport_crank_nicolson.nc differ diff --git a/torax/tests/test_data/test_all_transport_crank_nicolson.py b/torax/tests/test_data/test_all_transport_crank_nicolson.py index e2016c3a..f403aff7 100644 --- a/torax/tests/test_data/test_all_transport_crank_nicolson.py +++ b/torax/tests/test_data/test_all_transport_crank_nicolson.py @@ -38,7 +38,7 @@ 'current_eq': True, 'resistivity_mult': 100, # to shorten current diffusion time 't_final': 2, - 'largeValue_n': 1.0e5, + 'largeValue_n': 3.0e5, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_bremsstrahlung.nc b/torax/tests/test_data/test_bremsstrahlung.nc index 0d4c95b6..c270126c 100644 Binary files a/torax/tests/test_data/test_bremsstrahlung.nc and b/torax/tests/test_data/test_bremsstrahlung.nc differ diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc index 7277a35d..d297ee9e 100644 Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ diff --git a/torax/tests/test_data/test_changing_config_after.py b/torax/tests/test_data/test_changing_config_after.py index 131fffe1..a9222e4a 100644 --- a/torax/tests/test_data/test_changing_config_after.py +++ b/torax/tests/test_data/test_changing_config_after.py @@ -82,12 +82,6 @@ # likely be increased further beyond this default. 'dtmult': 30, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc index 7d7a3330..527f4cf4 100644 Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ diff --git a/torax/tests/test_data/test_changing_config_before.py b/torax/tests/test_data/test_changing_config_before.py index adb6f8a8..6d97f024 100644 --- a/torax/tests/test_data/test_changing_config_before.py +++ b/torax/tests/test_data/test_changing_config_before.py @@ -83,12 +83,6 @@ # likely be increased further beyond this default. 'dtmult': 30, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_chease.nc b/torax/tests/test_data/test_chease.nc index 45cf7acf..3ad3f949 100644 Binary files a/torax/tests/test_data/test_chease.nc and b/torax/tests/test_data/test_chease.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc index b4d806a0..ecb841dd 100644 Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.py b/torax/tests/test_data/test_iterbaseline_mockup.py index 77299c0e..f5d3e58d 100644 --- a/torax/tests/test_data/test_iterbaseline_mockup.py +++ b/torax/tests/test_data/test_iterbaseline_mockup.py @@ -67,12 +67,6 @@ 'el_heat_eq': True, 'current_eq': True, 'dens_eq': True, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc index 42caad0a..7baa348c 100644 Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_mockup.py b/torax/tests/test_data/test_iterhybrid_mockup.py index f82a0c29..2f805bcc 100644 --- a/torax/tests/test_data/test_iterhybrid_mockup.py +++ b/torax/tests/test_data/test_iterhybrid_mockup.py @@ -68,12 +68,6 @@ # likely be increased further beyond this default. 'dtmult': 50, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_iterhybrid_newton.nc b/torax/tests/test_data/test_iterhybrid_newton.nc index de73cbcd..3a089e49 100644 Binary files a/torax/tests/test_data/test_iterhybrid_newton.nc and b/torax/tests/test_data/test_iterhybrid_newton.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_newton.py b/torax/tests/test_data/test_iterhybrid_newton.py index 9ae8f120..aed6037a 100644 --- a/torax/tests/test_data/test_iterhybrid_newton.py +++ b/torax/tests/test_data/test_iterhybrid_newton.py @@ -72,12 +72,6 @@ # likely be increased further beyond this default. 'dtmult': 30, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary\ - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc index b9de2fa6..f6d4beaa 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.py b/torax/tests/test_data/test_iterhybrid_predictor_corrector.py index 3a3c3338..5331feef 100644 --- a/torax/tests/test_data/test_iterhybrid_predictor_corrector.py +++ b/torax/tests/test_data/test_iterhybrid_predictor_corrector.py @@ -68,12 +68,6 @@ # likely be increased further beyond this default. 'dtmult': 50, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary\ - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc index 1fb558d9..527f4cf4 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup.py b/torax/tests/test_data/test_iterhybrid_rampup.py index 5ee04b83..a639187d 100644 --- a/torax/tests/test_data/test_iterhybrid_rampup.py +++ b/torax/tests/test_data/test_iterhybrid_rampup.py @@ -77,12 +77,6 @@ # likely be increased further beyond this default. 'dtmult': 30, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_iterhybrid_rampup_short.py b/torax/tests/test_data/test_iterhybrid_rampup_short.py index 1d98bdef..a243674c 100644 --- a/torax/tests/test_data/test_iterhybrid_rampup_short.py +++ b/torax/tests/test_data/test_iterhybrid_rampup_short.py @@ -76,12 +76,6 @@ 'dens_eq': True, 'maxdt': 0.5, 'dt_reduction_factor': 3, - # effective source to dominate PDE in internal boundary condtion - # location if T != Tped - 'largeValue_T': 1.0e10, - # effective source to dominate density PDE in internal boundary - # condtion location if n != neped - 'largeValue_n': 1.0e8, }, }, 'geometry': { diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc index e0f1d3d2..f8f4554e 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc index f65a3d58..182cd87e 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_defromchie.nc and b/torax/tests/test_data/test_ne_qlknn_defromchie.nc differ diff --git a/torax/tests/test_data/test_ohmic_power.nc b/torax/tests/test_data/test_ohmic_power.nc index a6fba100..ff7d7e90 100644 Binary files a/torax/tests/test_data/test_ohmic_power.nc and b/torax/tests/test_data/test_ohmic_power.nc differ diff --git a/torax/tests/test_data/test_pc_method_ne.nc b/torax/tests/test_data/test_pc_method_ne.nc index a0a8cdde..6aec00c6 100644 Binary files a/torax/tests/test_data/test_pc_method_ne.nc and b/torax/tests/test_data/test_pc_method_ne.nc differ diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc index a22e1f56..fa98afa6 100644 Binary files a/torax/tests/test_data/test_prescribed_timedependent_ne.nc and b/torax/tests/test_data/test_prescribed_timedependent_ne.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_chease.nc b/torax/tests/test_data/test_psichease_ip_chease.nc index a77f9e15..538bf106 100644 Binary files a/torax/tests/test_data/test_psichease_ip_chease.nc and b/torax/tests/test_data/test_psichease_ip_chease.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_parameters.nc b/torax/tests/test_data/test_psichease_ip_parameters.nc index 23fecb3f..acfc7d30 100644 Binary files a/torax/tests/test_data/test_psichease_ip_parameters.nc and b/torax/tests/test_data/test_psichease_ip_parameters.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_johm.nc b/torax/tests/test_data/test_psichease_prescribed_johm.nc index e3c85e8c..6973555f 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_johm.nc and b/torax/tests/test_data/test_psichease_prescribed_johm.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot.nc b/torax/tests/test_data/test_psichease_prescribed_jtot.nc index cb1a61d3..241cef58 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot.nc differ diff --git a/torax/tests/test_data/test_qei_chease_highdens.nc b/torax/tests/test_data/test_qei_chease_highdens.nc index 3ffa4a0c..267fed0f 100644 Binary files a/torax/tests/test_data/test_qei_chease_highdens.nc and b/torax/tests/test_data/test_qei_chease_highdens.nc differ diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc index 5f892218..607366c3 100644 Binary files a/torax/tests/test_data/test_timedependence.nc and b/torax/tests/test_data/test_timedependence.nc differ diff --git a/torax/tests/test_lib/explicit_stepper.py b/torax/tests/test_lib/explicit_stepper.py index 4906c9c7..1ee6cfa4 100644 --- a/torax/tests/test_lib/explicit_stepper.py +++ b/torax/tests/test_lib/explicit_stepper.py @@ -99,7 +99,6 @@ def __call__( * true_ni_face * consts.keV2J * dynamic_runtime_params_slice_t.transport.chii_const - / geo_t.rmax**2 ) c_mat, c = diffusion_terms.make_diffusion_terms( diff --git a/torax/tests/test_lib/torax_refs.py b/torax/tests/test_lib/torax_refs.py index 88d31f77..6f8c00fc 100644 --- a/torax/tests/test_lib/torax_refs.py +++ b/torax/tests/test_lib/torax_refs.py @@ -61,7 +61,8 @@ def build_consistent_dynamic_runtime_params_slice_and_geometry( runtime_params, transport_getter=lambda: None, sources_getter=lambda: sources, - stepper_getter=lambda: None), + stepper_getter=lambda: None, + ), geometry_provider, ) @@ -248,118 +249,118 @@ def chease_references_Ip_from_chease() -> References: # pylint: disable=invalid psi = fvm.cell_variable.CellVariable( value=jnp.array( np.array([ - '0.028269200893290054', - '0.25830070956630535', - '0.7515864371427805', - '1.5290377293407753', - '2.6160416591477507', - '4.109776781076197', - '6.083267597177049', - '8.548498997963945', - '11.418026005498554', - '14.54591564529029', - '17.802294247322177', - '21.11023166707459', - '24.426030866786014', - '27.71745280935002', - '30.954224195485104', - '34.112857792036955', - '37.173258956533324', - '40.12077464958477', - '42.93919885653815', - '45.61832569488219', - '48.144560261511515', - '50.492447775755906', - '52.668083031356424', - '54.76130159148528', - '56.79525122154844', + 2.82691562998223e-02, + 2.58294982485687e-01, + 7.51532731064005e-01, + 1.52898638812398e00, + 2.61597539330777e00, + 4.10955734212211e00, + 6.08340992472874e00, + 8.54828085960123e00, + 1.14180751339249e01, + 1.45448690085190e01, + 1.78016419395054e01, + 2.11095105087725e01, + 2.44259250193576e01, + 2.77169841485157e01, + 3.09538089854500e01, + 3.41121726512746e01, + 3.71726376565201e01, + 4.01195454787118e01, + 4.29382385995525e01, + 4.56175767225333e01, + 4.81440970724439e01, + 5.04925801508127e01, + 5.26679900645839e01, + 5.47613521745655e01, + 5.67944467049695e01, ]).astype('float64') ), right_face_grad_constraint=jnp.array(50.417485084359726), dr=geo.dr_norm, ) psi_face_grad = np.array([ - '0.0', - '5.750787716825383', - '12.332143189411878', - '19.43628230494987', - '27.175098245174382', - '37.34337804821115', - '49.33727040252131', - '61.63078501967238', - '71.73817518836523', - '78.19724099479339', - '81.40946505079718', - '82.69843549381032', - '82.89497999278561', - '82.28554856410017', - '80.91928465337705', - '78.96583991379629', - '76.51002911240923', - '73.68789232628608', - '70.46060517383452', - '66.97817095860098', - '63.15586416573318', - '58.69718785610978', - '54.39088139001296', - '52.330464003221344', - '50.84874075157906', - '50.417485084359726', + 0.0, + 5.75064565464661, + 12.33094371445794, + 19.43634142649927, + 27.17472512959477, + 37.33954872035863, + 49.34631456516576, + 61.62177337181222, + 71.7448568580918, + 78.16984686485159, + 81.41932327466112, + 82.69671423167608, + 82.91036276462754, + 82.27647822895267, + 80.92062092335786, + 78.95909164561567, + 76.51162513113849, + 73.67269555479083, + 70.46732802101872, + 66.98345307451862, + 63.16300874776495, + 58.71207695922127, + 54.38524784427834, + 52.33405274954137, + 50.82736326009893, + 50.41748508435973, ]).astype('float64') jtot = np.array([ - 795930.385028747, - 840806.1886064387, - 907353.8327698484, - 997244.898488572, - 1147270.446420881, - 1299395.796373129, - 1355791.8832536936, - 1271027.2034415167, - 1082414.935365022, - 883539.1967190297, - 735591.697040288, - 636536.3331668673, - 560950.4569715615, - 495141.107022868, - 438029.14468179113, - 390516.4836237442, - 349981.38558092073, - 315240.739563484, - 286627.17652595515, - 257638.8196543309, - 242652.4683838375, - 282852.536672287, - 395751.3390132676, - 521040.3581048132, - 579217.4375589897, + 795917.9083302062, + 840782.4499447887, + 907361.0199350917, + 997214.1844258476, + 1147368.4475644985, + 1299436.7025480503, + 1355582.6447244585, + 1270699.3449932635, + 1082130.7974471736, + 883946.832547368, + 736175.4601967924, + 636551.6005556053, + 560567.6935966521, + 494898.79834195157, + 438073.51598169986, + 390326.22609503136, + 349883.16942411405, + 315812.03175793774, + 287118.2725221936, + 257945.50531862015, + 242596.7573229568, + 282222.4593829316, + 394854.20010322006, + 520292.415208095, + 580718.6763027112, ]) s = np.array([ - -0.07221339029182716, - -0.07221339029182716, - -0.10978697094216668, - -0.14556469833048682, - -0.3182545587996119, - -0.48410241240691904, - -0.4777827504007501, - -0.27329938411328436, - 0.07676160004898383, - 0.4477451568703172, - 0.7339849318492895, - 0.9180790330929585, - 1.056015149635723, - 1.1952782347275002, - 1.3428889005376405, - 1.4969801606698092, - 1.6477614361846962, - 1.822310073007422, - 2.0162016336630817, - 2.2346589515146733, - 2.570972907161161, - 2.897484271282114, - 2.626089343552705, - 2.078461443212857, - 1.7481696293083797, - 1.8139935034456611, + -0.07213558746176, + -0.07213558746176, + -0.10991003958013, + -0.14562624389212, + -0.31798201650368, + -0.48489144551329, + -0.47719517520208, + -0.27334155928272, + 0.07788007874047, + 0.44735990737644, + 0.73241268296836, + 0.917711631104, + 1.05655725949415, + 1.19644939698028, + 1.34268091118458, + 1.49699333857542, + 1.64867425366475, + 1.82182056987522, + 2.01325973255752, + 2.23449608843021, + 2.56911309068007, + 2.89960958458338, + 2.6288582770712, + 2.08242930227477, + 1.74940855765015, + 1.81507856025273, ]) return References( runtime_params=runtime_params, @@ -400,118 +401,118 @@ def chease_references_Ip_from_runtime_params() -> References: # pylint: disable psi = fvm.cell_variable.CellVariable( value=jnp.array( np.array([ - '0.03602782820394983', - '0.32919266534421404', - '0.9578631932332422', - '1.9486899837739167', - '3.334027722464634', - '5.2377261169890375', - '7.752851618871569', - '10.894678433354123', - '14.551761859387222', - '18.538116868455745', - '22.688225295074556', - '26.904043121624614', - '31.12988043400343', - '35.324650025865274', - '39.449769935526916', - '43.475306738118505', - '47.3756507133815', - '51.13212722004177', - '54.72408241952786', - '58.138509372479355', - '61.35808198484154', - '64.35035928071211', - '67.12310880126317', - '69.79082194116809', - '72.38299948888208', + 3.60277713715760e-02, + 3.29185366436919e-01, + 9.57794747245436e-01, + 1.94862455169660e00, + 3.33394326962468e00, + 5.23744645188747e00, + 7.75303300895049e00, + 1.08944004257979e01, + 1.45518244713624e01, + 1.85367829768494e01, + 2.26873939580287e01, + 2.69031240377243e01, + 3.11297455362048e01, + 3.53240527386590e01, + 3.94492407689661e01, + 4.34744335569566e01, + 4.73748588943901e01, + 5.11305606969125e01, + 5.47228586150829e01, + 5.81375548408243e01, + 6.13574916711564e01, + 6.43505279867828e01, + 6.71229903192826e01, + 6.97908864070344e01, + 7.23819741685979e01, ]).astype('float64') ), right_face_grad_constraint=jnp.array(64.25482269382654), dr=geo.dr_norm, ) psi_face_grad = np.array([ - '0.0', - '7.329120928506605', - '15.716763197225703', - '24.77066976351686', - '34.63344346726794', - '47.59245986311008', - '62.87813754706328', - '78.54567036206386', - '91.42708565082746', - '99.65887522671308', - '103.75271066547027', - '105.39544566375145', - '105.64593280947037', - '104.86923979654614', - '103.12799774154104', - '100.63842006478971', - '97.50859938157497', - '93.91191266650658', - '89.79887998715235', - '85.36067382378735', - '80.48931530905463', - '74.8069323967643', - '69.31873801377648', - '66.69282849762297', - '64.80443869284969', - '64.25482269382654', + 0.0, + 7.32893987663357, + 15.71523452021292, + 24.77074511127902, + 34.63296794820204, + 47.58757955656982, + 62.88966392657549, + 78.5341854211854, + 91.4356011391122, + 99.62396263717564, + 103.76527452948281, + 105.39325199238804, + 105.66553746201474, + 104.85768006135298, + 103.12970075767804, + 100.62981969976352, + 97.51063343583776, + 93.89254506305899, + 89.80744795425916, + 85.36740564353558, + 80.49842075830291, + 74.82590789066083, + 69.31155831249427, + 66.69740219379392, + 64.7771940390875, + 64.25482269382653, ]).astype('float64') jtot = np.array([ - 1014377.5652648797, - 1071569.7635132223, - 1156381.750252685, - 1270943.8804313582, - 1462144.7101793122, - 1656021.643391323, - 1727895.9258366497, - 1619867.1444938849, - 1379489.2711656648, - 1126031.0650806061, - 937478.6145974831, - 811236.9975888315, - 714906.1896367776, - 631035.1257494364, - 558248.4921483698, - 497695.7373473658, - 446035.5735518052, - 401760.1788869469, - 365293.47658042004, - 328349.1163479699, - 309249.6840358926, - 360482.8674410441, - 504367.3256727074, - 664042.5592497136, - 738186.6368235692, + 1014361.6642723732, + 1071539.5096540153, + 1156390.909972325, + 1270904.7367364431, + 1462269.608239906, + 1656073.7764759487, + 1727629.2607189012, + 1619449.3036113628, + 1379127.150137565, + 1126550.5786547544, + 938222.5945491319, + 811256.4551907588, + 714418.374888706, + 630726.3142071059, + 558305.0413793066, + 497453.2626631715, + 445910.40146656823, + 402488.2651570188, + 365919.3563938611, + 328739.97346724605, + 309178.6827883966, + 359679.8622050641, + 503223.9623832343, + 663089.3395851713, + 740099.8982473996, ]) s = np.array([ - -0.07221339029182716, - -0.07221339029182716, - -0.10978697094216645, - -0.1455646983304877, - -0.3182545587996117, - -0.4841024124069177, - -0.47778275040075163, - -0.2732993841132834, - 0.07676160004898405, - 0.447745156870319, - 0.7339849318492886, - 0.9180790330929495, - 1.0560151496357255, - 1.1952782347274975, - 1.3428889005376485, - 1.4969801606698174, - 1.647761436184692, - 1.822310073007411, - 2.0162016336630986, - 2.2346589515146715, - 2.570972907161099, - 2.8974842712821975, - 2.626089343552712, - 2.0784614432128303, - 1.7481696293084164, - 1.813993503445698, + -0.07213558746176, + -0.07213558746176, + -0.10991003958012, + -0.14562624389212, + -0.31798201650368, + -0.48489144551329, + -0.47719517520207, + -0.27334155928274, + 0.07788007874055, + 0.4473599073764, + 0.73241268296826, + 0.91771163110416, + 1.05655725949404, + 1.19644939698014, + 1.34268091118474, + 1.49699333857536, + 1.64867425366483, + 1.82182056987507, + 2.01325973255764, + 2.2344960884306, + 2.56911309067965, + 2.89960958458354, + 2.62885827707079, + 2.0824293022747, + 1.74940855765083, + 1.81507856025345, ]) return References( runtime_params=runtime_params, diff --git a/torax/tests/test_run_simulation_main.py b/torax/tests/test_run_simulation_main.py index 6030d0d9..1e5f9b64 100644 --- a/torax/tests/test_run_simulation_main.py +++ b/torax/tests/test_run_simulation_main.py @@ -231,7 +231,7 @@ def check(output_path, ground_truth_path): # the mode=zero case, needing looser tolerance for this than # for other tests. # rtol=0.0, - atol=1.0e-9, + atol=5.0e-5, ): diff = ov - gv max_diff = np.abs(diff).max() diff --git a/torax/time_step_calculator/chi_time_step_calculator.py b/torax/time_step_calculator/chi_time_step_calculator.py index a2df6224..0dd29d80 100644 --- a/torax/time_step_calculator/chi_time_step_calculator.py +++ b/torax/time_step_calculator/chi_time_step_calculator.py @@ -81,7 +81,7 @@ def next_dt( chi_max = core_transport.chi_max(geo) - basic_dt = (3.0 / 4.0) * (geo.dr_norm**2) / chi_max * geo.rmax**2 + basic_dt = (3.0 / 4.0) * (geo.dr_norm**2) / chi_max dt = jnp.minimum( dynamic_runtime_params_slice.numerics.dtmult * basic_dt, diff --git a/torax/transport_model/qlknn_wrapper.py b/torax/transport_model/qlknn_wrapper.py index e9ba33a0..af808145 100644 --- a/torax/transport_model/qlknn_wrapper.py +++ b/torax/transport_model/qlknn_wrapper.py @@ -404,11 +404,14 @@ def make_core_transport( # For small density gradients or up-gradient transport, set pure effective # convection. Otherwise pure effective diffusion. def DVeff_approach() -> tuple[jax.Array, jax.Array]: + # The geo.rmax is to unnormalize the face_grad. Deff = -pfe_SI / ( - core_profiles.ne.face_grad() * geo.g1_over_vpr2_face / geo.rmax + core_profiles.ne.face_grad() * geo.g1_over_vpr2_face * geo.rmax + constants.eps ) - Veff = pfe_SI / (core_profiles.ne.face_value() * geo.g0_over_vpr_face) + Veff = pfe_SI / ( + core_profiles.ne.face_value() * geo.g0_over_vpr_face * geo.rmax + ) Deff_mask = ( ((pfe >= 0) & (prepared_data['Ane'] >= 0)) | ((pfe < 0) & (prepared_data['Ane'] < 0)) @@ -432,7 +435,8 @@ def Dscaled_approach() -> tuple[jax.Array, jax.Array]: * d_face_el / prepared_data['Rmaj'] * geo.g1_over_vpr2_face - ) / geo.g0_over_vpr_face + * geo.rmax**2 + ) / (geo.g0_over_vpr_face * geo.rmax) return d_face_el, v_face_el d_face_el, v_face_el = jax.lax.cond( @@ -585,9 +589,7 @@ def _combined( ) def __hash__(self) -> int: - return hash( - ('QLKNNTransportModel' + self._model_path) - ) + return hash(('QLKNNTransportModel' + self._model_path)) def __eq__(self, other: QLKNNTransportModel) -> bool: return (