Skip to content

Commit 5681d7a

Browse files
theo-browntheo-brown
committed
Update units and normalisations
1 parent 3ec4237 commit 5681d7a

File tree

1 file changed

+166
-159
lines changed

1 file changed

+166
-159
lines changed

‎torax/transport_model/tglf_based_transport_model.py

+166-159
Original file line numberDiff line numberDiff line change
@@ -26,179 +26,186 @@
2626

2727
@chex.dataclass
2828
class RuntimeParams(quasilinear_transport_model.RuntimeParams):
29-
"""Shared parameters for TGLF-based models."""
29+
"""Shared parameters for TGLF-based models."""
3030

31-
def make_provider(
32-
self, torax_mesh: geometry.Grid1D | None = None
33-
) -> "RuntimeParamsProvider":
34-
return RuntimeParamsProvider(**self.get_provider_kwargs(torax_mesh))
31+
def make_provider(
32+
self, torax_mesh: geometry.Grid1D | None = None
33+
) -> "RuntimeParamsProvider":
34+
return RuntimeParamsProvider(**self.get_provider_kwargs(torax_mesh))
3535

3636

3737
@chex.dataclass(frozen=True)
3838
class DynamicRuntimeParams(quasilinear_transport_model.DynamicRuntimeParams):
39-
"""Shared parameters for TGLF-based models."""
39+
"""Shared parameters for TGLF-based models."""
4040

41-
pass
41+
pass
4242

4343

4444
@chex.dataclass
4545
class RuntimeParamsProvider(runtime_params_lib.RuntimeParamsProvider):
46-
"""Provides a RuntimeParams to use during time t of the sim."""
46+
"""Provides a RuntimeParams to use during time t of the sim."""
4747

48-
runtime_params_config: RuntimeParams
48+
runtime_params_config: RuntimeParams
4949

50-
def build_dynamic_params(self, t: chex.Numeric) -> DynamicRuntimeParams:
51-
return DynamicRuntimeParams(**self.get_dynamic_params_kwargs(t))
50+
def build_dynamic_params(self, t: chex.Numeric) -> DynamicRuntimeParams:
51+
return DynamicRuntimeParams(**self.get_dynamic_params_kwargs(t))
5252

5353

5454
@chex.dataclass(frozen=True)
5555
class TGLFInputs(quasilinear_transport_model.QuasilinearInputs):
56-
r"""Dimensionless inputs to TGLF-based models.
57-
58-
See https://gafusion.github.io/doc/tglf/tglf_table.html for definitions.
59-
"""
60-
61-
# Ti/Te
62-
Ti_over_Te: chex.Array
63-
# dRmaj/dr
64-
dRmaj: chex.Array
65-
# q
66-
q: chex.Array
67-
# r/q dq/dr
68-
s_hat: chex.Array
69-
# nu_ee (see note in prepare_tglf_inputs)
70-
nu_ee: chex.Array
71-
# Elongation, kappa
72-
kappa: chex.Array
73-
# Shear in elongation, r/kappa dkappa/dr
74-
kappa_shear: chex.Array
75-
# Triangularity, delta
76-
delta: chex.Array
77-
# Shear in triangularity, r ddelta/dr
78-
delta_shear: chex.Array
79-
# Electron pressure defined w.r.t B_unit
80-
beta_e: chex.Array
81-
# Effective charge
82-
Zeff: chex.Array
83-
84-
85-
class TGLFBasedTransportModel(quasilinear_transport_model.QuasilinearTransportModel):
86-
"""Base class for TGLF-based transport models."""
87-
88-
def _prepare_tglf_inputs(
89-
Zeff_face: chex.Array,
90-
q_correction_factor: chex.Numeric,
91-
geo: geometry.Geometry,
92-
core_profiles: state.CoreProfiles,
93-
) -> TGLFInputs:
94-
## Shorthand 'standard' variables
95-
Te = core_profiles.temp_el
96-
Ti = core_profiles.temp_ion
97-
ne = core_profiles.ne * core_profiles.nref
98-
# q must be recalculated since in the nonlinear solver psi has intermediate
99-
# states in the iterative solve
100-
q, _ = physics.calc_q_from_psi(
101-
geo=geo,
102-
psi=core_profiles.psi,
103-
q_correction_factor=q_correction_factor,
104-
)
105-
106-
## Reference values used for TGLF-specific normalisation
107-
# https://gafusion.github.io/doc/cgyro/outputs.html#output-normalization
108-
# Note: B_unit = q/r dpsi/dr [SOURCE?]
109-
# Note: TGLF uses geo.rmid = (Rmax - Rmin)/2 as the radial coordinate
110-
# This means all gradients are calculated w.r.t. rmid
111-
vref = (Te.face_value() / (core_profiles.Ai * CONSTANTS.mp)) ** 0.5
112-
lref = geo.Rmin[-1] # Minor radius at LCFS
113-
B_unit = (
114-
q / geo.rmid * jnp.gradient(core_profiles.psi, geo.rmid)
115-
)
116-
117-
## Dimensionless gradients, eg lref_over_lti where lref=amin, lti = -ti / (dti/dr)
118-
normalized_log_gradients = quasilinear_transport_model.NormalizedLogarithmicGradients.from_profiles(
119-
core_profiles=core_profiles,
120-
radial_coordinate=geo.rmid,
121-
reference_length=lref,
122-
)
123-
124-
## Dimensionless temperature ratio
125-
Ti_over_Te = Ti.face_value() / Te.face_value()
126-
127-
## Dimensionless electron-electron collision frequency = nu_ee / (vref/lref)
128-
# https://gafusion.github.io/doc/tglf/tglf_list.html#xnue
129-
# https://gafusion.github.io/doc/cgyro/cgyro_list.html#cgyro-nu-ee
130-
# Note: In the TGLF docs, XNUE is mislabelled as electron-ion collision frequency.
131-
# It is actually the electron-electron collision frequency, and is defined as in CGYRO
132-
# See https://pyrokinetics.readthedocs.io/en/latest/user_guide/collisions.html#tglf
133-
Lambda_ee = physics._calculate_lambda_ee(Te, ne)
134-
normalised_nu_ee = (4 * jnp.pi * ne * CONSTANTS.qe**4 * Lambda_ee) / (
135-
CONSTANTS.me**0.5 * (2 * Te) ** 1.5
136-
)
137-
nu_ee = normalised_nu_ee / (vref / lref)
138-
139-
## Safety factor, q
140-
# https://gafusion.github.io/doc/tglf/tglf_list.html#q-sa
141-
# defined before
142-
143-
## Safety factor shear, s_hat = r/q dq/dr
144-
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-shat-sa
145-
# Note: calc_s_from_psi_rmid gives rq dq/dr, so we divide by q**2
146-
s_hat = physics.calc_s_from_psi_rmid(geo, core_profiles.psi) / q**2
147-
148-
## Electron beta
149-
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-betae
150-
p_e = ne * (Te * 1e3) # ne in m^-3, Te in eV
151-
beta_e = 8 * jnp.pi * p_e / B_unit**2
152-
153-
## Major radius shear = dRmaj/dr
154-
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-drmajdx-loc
155-
dRmaj = jnp.gradient(geo.Rmaj, geo.rmid)
156-
157-
## Elongation shear = r/kappa dkappa/dr
158-
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-s-kappa-loc
159-
kappa = geo.elongation_face
160-
kappa_shear = geo.rmid_face / kappa * jnp.gradient(kappa, geo.rmid_face)
161-
162-
## Triangularity shear = r ddelta/dr
163-
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-s-delta-loc
164-
delta = geo.delta_face
165-
delta_shear = geo.rmid_face * jnp.gradient(delta, geo.rmid_face)
166-
167-
## Gyrobohm diffusivity
168-
# https://gafusion.github.io/doc/tglf/tglf_table.html#id7
169-
# https://gafusion.github.io/doc/cgyro/outputs.html#output-normalization
170-
# Note: TGLF uses the same normalisation as CGYRO, ie
171-
# chiGB = ne * vref * T_e * (rho_s_unit / lref)**2
172-
# where rho_s_unit = vref / (e*B_unit/ m_D / c)
173-
# TODO: Check if this code implementation is correct/equivalent to the above
174-
chiGB = quasilinear_transport_model.calculate_chiGB(
175-
reference_temperature=core_profiles.temp_ion.face_value(),
176-
reference_magnetic_field=geo.B0,
177-
reference_mass=core_profiles.Ai,
178-
reference_length=geo.Rmin,
179-
)
180-
# gammaGB =
181-
182-
return TGLFInputs(
183-
# From QuasilinearInputs
184-
chiGB=chiGB,
185-
Rmin=geo.Rmin,
186-
Rmaj=geo.Rmaj,
187-
lref_over_lti=normalized_log_gradients.lref_over_lti,
188-
lref_over_lte=normalized_log_gradients.lref_over_lte,
189-
lref_over_lne=normalized_log_gradients.lref_over_lne,
190-
lref_over_lni0=normalized_log_gradients.lref_over_lni0,
191-
lref_over_lni1=normalized_log_gradients.lref_over_lni1,
192-
# From TGLFInputs
193-
Ti_over_Te=Ti_over_Te,
194-
dRmaj=dRmaj,
195-
q=q,
196-
s_hat=s_hat,
197-
nu_ee=nu_ee,
198-
kappa=kappa,
199-
kappa_shear=kappa_shear,
200-
delta=delta,
201-
delta_shear=delta_shear,
202-
beta_e=beta_e,
203-
Zeff=Zeff_face,
56+
r"""Dimensionless inputs to TGLF-based models.
57+
58+
See https://gafusion.github.io/doc/tglf/tglf_table.html for definitions.
59+
"""
60+
61+
# Ti/Te
62+
Ti_over_Te: chex.Array
63+
# dRmaj/dr
64+
dRmaj: chex.Array
65+
# q
66+
q: chex.Array
67+
# r/q dq/dr
68+
s_hat: chex.Array
69+
# nu_ee (see note in prepare_tglf_inputs)
70+
nu_ee: chex.Array
71+
# Elongation, kappa
72+
kappa: chex.Array
73+
# Shear in elongation, r/kappa dkappa/dr
74+
kappa_shear: chex.Array
75+
# Triangularity, delta
76+
delta: chex.Array
77+
# Shear in triangularity, r ddelta/dr
78+
delta_shear: chex.Array
79+
# Electron pressure defined w.r.t B_unit
80+
beta_e: chex.Array
81+
# Effective charge
82+
Zeff: chex.Array
83+
84+
85+
class TGLFBasedTransportModel(
86+
quasilinear_transport_model.QuasilinearTransportModel
87+
):
88+
"""Base class for TGLF-based transport models."""
89+
90+
def _prepare_tglf_inputs(
91+
Zeff_face: chex.Array,
92+
q_correction_factor: chex.Numeric,
93+
geo: geometry.Geometry,
94+
core_profiles: state.CoreProfiles,
95+
) -> TGLFInputs:
96+
## Shorthand 'standard' variables
97+
Te_keV = core_profiles.temp_el.face_value()
98+
Te_eV = Te_keV * 1e3
99+
Te_J = Te_keV * CONSTANTS.keV2J
100+
Ti_keV = core_profiles.temp_ion.face_value()
101+
ne = core_profiles.ne * core_profiles.nref
102+
# q must be recalculated since in the nonlinear solver psi has intermediate
103+
# states in the iterative solve
104+
q, _ = physics.calc_q_from_psi(
105+
geo=geo,
106+
psi=core_profiles.psi,
107+
q_correction_factor=q_correction_factor,
108+
)
109+
110+
## Reference values used for TGLF-specific normalisation
111+
# https://gafusion.github.io/doc/cgyro/outputs.html#output-normalization
112+
# https://gafusion.github.io/doc/geometry.html#effective-field
113+
# B_unit = 1/r d(psi_tor)/dr = q/r dpsi/dr
114+
# Note: TGLF uses geo.rmid = (Rmax - Rmin)/2 as the radial coordinate
115+
# This means all gradients are calculated w.r.t. rmid
116+
m_D_amu = 2.014 # Mass of deuterium
117+
m_D = m_D_amu * CONSTANTS.mp # Mass of deuterium
118+
c_s = (Te_J / m_D) ** 0.5
119+
a = geo.Rmin[-1] # Minor radius at LCFS
120+
B_unit = q / (geo.rmid) * jnp.gradient(core_profiles.psi, geo.rmid)
121+
122+
## Dimensionless gradients, eg lref_over_lti where lref=amin, lti = -ti / (dti/dr)
123+
normalized_log_gradients = quasilinear_transport_model.NormalizedLogarithmicGradients.from_profiles(
124+
core_profiles=core_profiles,
125+
radial_coordinate=geo.rmid,
126+
reference_length=a,
127+
)
128+
129+
## Dimensionless temperature ratio
130+
Ti_over_Te = Ti_keV / Te_keV
131+
132+
## Dimensionless electron-electron collision frequency = nu_ee / (c_s/a)
133+
# https://gafusion.github.io/doc/tglf/tglf_list.html#xnue
134+
# https://gafusion.github.io/doc/cgyro/cgyro_list.html#cgyro-nu-ee
135+
# Note: In the TGLF docs, XNUE is mislabelled as electron-ion collision frequency.
136+
# It is actually the electron-electron collision frequency, and is defined as in CGYRO
137+
# See https://pyrokinetics.readthedocs.io/en/latest/user_guide/collisions.html#tglf
138+
# Lambda_ee is computed with keV and m^-3 units
139+
# normalised_nu_ee is computed with SI units (ie J rather than keV)
140+
Lambda_ee = physics._calculate_lambda_ee(Te_keV, ne)
141+
normalised_nu_ee = (4 * jnp.pi * ne * CONSTANTS.qe**4 * Lambda_ee) / (
142+
CONSTANTS.me**0.5 * (2 * Te_J) ** 1.5
143+
)
144+
nu_ee = normalised_nu_ee / (c_s / a)
145+
146+
## Safety factor, q
147+
# https://gafusion.github.io/doc/tglf/tglf_list.html#q-sa
148+
# defined before
149+
150+
## Safety factor shear, s_hat = r/q dq/dr
151+
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-shat-sa
152+
# Note: calc_s_from_psi_rmid gives rq dq/dr, so we divide by q**2
153+
s_hat = physics.calc_s_from_psi_rmid(geo, core_profiles.psi) / q**2
154+
155+
## Electron beta
156+
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-betae
157+
# Note: Te in eV
158+
beta_e = 8 * jnp.pi * ne * Te_eV / B_unit**2
159+
160+
## Major radius shear = dRmaj/dr
161+
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-drmajdx-loc
162+
dRmaj = jnp.gradient(geo.Rmaj, geo.rmid)
163+
164+
## Elongation shear = r/kappa dkappa/dr
165+
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-s-kappa-loc
166+
kappa = geo.elongation_face
167+
kappa_shear = geo.rmid_face / kappa * jnp.gradient(kappa, geo.rmid_face)
168+
169+
## Triangularity shear = r ddelta/dr
170+
# https://gafusion.github.io/doc/tglf/tglf_list.html#tglf-s-delta-loc
171+
delta = geo.delta_face
172+
delta_shear = geo.rmid_face * jnp.gradient(delta, geo.rmid_face)
173+
174+
## Gyrobohm diffusivity
175+
# https://gafusion.github.io/doc/tglf/tglf_table.html#id7
176+
# https://gafusion.github.io/doc/cgyro/outputs.html#output-normalization
177+
# Note: TGLF uses the same normalisation as CGYRO
178+
# This has an extra c^2 factor compared to TORAX's calculate_chiGB
179+
chiGB = (
180+
quasilinear_transport_model.calculate_chiGB(
181+
reference_temperature=Te_keV, # conversion to J done internally
182+
reference_magnetic_field=B_unit,
183+
reference_mass=m_D_amu,
184+
reference_length=a,
204185
)
186+
* CONSTANTS.c**2
187+
)
188+
189+
return TGLFInputs(
190+
# From QuasilinearInputs
191+
chiGB=chiGB,
192+
Rmin=geo.Rmin,
193+
Rmaj=geo.Rmaj,
194+
lref_over_lti=normalized_log_gradients.lref_over_lti,
195+
lref_over_lte=normalized_log_gradients.lref_over_lte,
196+
lref_over_lne=normalized_log_gradients.lref_over_lne,
197+
lref_over_lni0=normalized_log_gradients.lref_over_lni0,
198+
lref_over_lni1=normalized_log_gradients.lref_over_lni1,
199+
# From TGLFInputs
200+
Ti_over_Te=Ti_over_Te,
201+
dRmaj=dRmaj,
202+
q=q,
203+
s_hat=s_hat,
204+
nu_ee=nu_ee,
205+
kappa=kappa,
206+
kappa_shear=kappa_shear,
207+
delta=delta,
208+
delta_shear=delta_shear,
209+
beta_e=beta_e,
210+
Zeff=Zeff_face,
211+
)

0 commit comments

Comments
 (0)