Skip to content

Commit b7bec88

Browse files
MateoBellMateoBell
committed
modify _load_imas_data to load any ids
1 parent 2606f4d commit b7bec88

1 file changed

Lines changed: 191 additions & 0 deletions

File tree

‎torax/_src/geometry/imas.py‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
# Copyright 2024 DeepMind Technologies Limited
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Useful functions for handling of IMAS IDSs."""
16+
import os
17+
from typing import Any
18+
19+
import imas
20+
from imas import ids_toplevel
21+
import numpy as np
22+
import scipy
23+
from torax._src.geometry import geometry_loader
24+
25+
26+
# pylint: disable=invalid-name
27+
def geometry_from_IMAS(
28+
geometry_directory: str | None = None,
29+
Ip_from_parameters: bool = False,
30+
n_rho: int = 25,
31+
hires_factor: int = 4,
32+
equilibrium_object: ids_toplevel.IDSToplevel | None = None,
33+
imas_uri: str | None = None,
34+
imas_filepath: str | None = None,
35+
) -> dict[str, Any]:
36+
"""Constructs a StandardGeometryIntermediates from a IMAS equilibrium IDS.
37+
38+
Currently written for COCOSv17 and DDv4.
39+
40+
Args:
41+
geometry_directory: Directory where to find the equilibrium object. If None,
42+
it defaults to another dir. See `load_geo_data` implementation.
43+
Ip_from_parameters: If True, the Ip is taken from the parameters and the
44+
values in the Geometry are resacled to match the new Ip.
45+
n_rho: Radial grid points (num cells)
46+
hires_factor: Grid refinement factor for poloidal flux <--> plasma current
47+
calculations.
48+
equilibrium_object: The equilibrium IDS containing the relevant data.
49+
imas_uri: The IMAS uri containing the equilibrium data.
50+
imas_filepath: The path to the IMAS netCDF file containing the equilibrium
51+
data.
52+
53+
Returns:
54+
A StandardGeometry instance based on the input file. This can then be
55+
used to build a StandardGeometry by passing to `build_standard_geometry`.
56+
"""
57+
# If the equilibrium_object is the file name, load the ids from the netCDF.
58+
if equilibrium_object is not None:
59+
equilibrium = equilibrium_object
60+
elif imas_uri is not None:
61+
equilibrium = _load_imas_data(
62+
imas_uri,
63+
geometry_directory,
64+
)
65+
elif imas_filepath is not None:
66+
equilibrium = _load_imas_data(
67+
imas_filepath,
68+
geometry_directory,
69+
)
70+
else:
71+
raise ValueError(
72+
"equilibrium_object must be a string (file path) or an IDS"
73+
)
74+
# TODO(b/431977390): Currently only the first time slice is used, extend to
75+
# support multiple time slices.
76+
IMAS_data = equilibrium.time_slice[0]
77+
B_0 = np.abs(IMAS_data.global_quantities.magnetic_axis.b_field_phi)
78+
R_major = np.asarray(equilibrium.vacuum_toroidal_field.r0)
79+
80+
# Poloidal flux.
81+
psi = 1 * IMAS_data.profiles_1d.psi # Sign changed ddv4
82+
83+
# Toroidal flux.
84+
phi = -1 * IMAS_data.profiles_1d.phi
85+
86+
# Midplane radii.
87+
R_in = IMAS_data.profiles_1d.r_inboard
88+
R_out = IMAS_data.profiles_1d.r_outboard
89+
# toroidal field flux function
90+
F = -1 * IMAS_data.profiles_1d.f
91+
92+
# Flux surface integrals of various geometry quantities.
93+
# IDS Contour integrals.
94+
if IMAS_data.profiles_1d.dvolume_dpsi:
95+
dvoldpsi = 1 * IMAS_data.profiles_1d.dvolume_dpsi # Sign changed ddv4.
96+
else:
97+
dvoldpsi = np.gradient(
98+
IMAS_data.profiles_1d.volume, IMAS_data.profiles_1d.psi
99+
)
100+
# dpsi_drho_tor
101+
if IMAS_data.profiles_1d.dpsi_drho_tor:
102+
dpsidrhotor = 1 * IMAS_data.profiles_1d.dpsi_drho_tor # Sign changed ddv4.
103+
else:
104+
rho_tor = IMAS_data.profiles_1d.rho_tor
105+
if not rho_tor:
106+
if B_0 is None or not IMAS_data.profiles_1d.phi:
107+
raise ValueError(
108+
"rho_tor not provided and cannot be calculated from given"
109+
" equilibrium IDS"
110+
)
111+
rho_tor = np.sqrt(IMAS_data.profiles_1d.phi / (np.pi * B_0))
112+
dpsidrhotor = np.gradient(IMAS_data.profiles_1d.psi, rho_tor)
113+
114+
flux_surf_avg_RBp = (
115+
IMAS_data.profiles_1d.gm7 * dpsidrhotor / (2 * np.pi)
116+
) # dpsi, C0/C1
117+
flux_surf_avg_R2Bp2 = (
118+
IMAS_data.profiles_1d.gm3 * (dpsidrhotor**2) / (4 * np.pi**2)
119+
) # C4/C1
120+
flux_surf_avg_Bp2 = (
121+
IMAS_data.profiles_1d.gm2 * (dpsidrhotor**2) / (4 * np.pi**2)
122+
) # C3/C1
123+
int_dl_over_Bp = dvoldpsi # C1
124+
flux_surf_avg_1_over_R2 = IMAS_data.profiles_1d.gm1 # C2/C1
125+
126+
# jtor = dI/drhon / (drho/dS) = dI/drhon / spr
127+
# spr = vpr / ( 2 * np.pi * R_major)
128+
# -> Ip_profile = integrate(y = spr * jtor, x= rhon, initial = 0.0)
129+
jtor = -1 * IMAS_data.profiles_1d.j_phi
130+
rhon = IMAS_data.profiles_1d.rho_tor_norm
131+
if not rhon:
132+
if B_0 is None or not IMAS_data.profiles_1d.phi:
133+
raise ValueError(
134+
"rho_tor_norm not provided and cannot be calculated from given"
135+
" equilibrium IDS"
136+
)
137+
rho_tor = np.sqrt(IMAS_data.profiles_1d.phi / (np.pi * B_0))
138+
rhon = rho_tor / rho_tor[-1]
139+
vpr = 4 * np.pi * phi[-1] * rhon / (F * flux_surf_avg_1_over_R2)
140+
spr = vpr / (2 * np.pi * R_major)
141+
# This Ip_profile by integration results in a discrepancy between this term
142+
# and the total Ip from IDS.
143+
Ip_profile_unscaled = scipy.integrate.cumulative_trapezoid(
144+
y=spr * jtor, x=rhon, initial=0.0
145+
)
146+
147+
# Because of the discrepancy between Ip_profile[-1] (computed by integration)
148+
# and global_quantities.ip, here we will scale Ip_profile such that the total
149+
# plasma current is equal.
150+
Ip_total = -1 * IMAS_data.global_quantities.ip
151+
Ip_profile = Ip_profile_unscaled * (Ip_total / Ip_profile_unscaled[-1])
152+
153+
z_magnetic_axis = np.asarray(IMAS_data.global_quantities.magnetic_axis.z)
154+
155+
return {
156+
"Ip_from_parameters": Ip_from_parameters,
157+
"R_major": R_major,
158+
"a_minor": np.asarray(IMAS_data.boundary.minor_radius),
159+
"B_0": B_0,
160+
"psi": psi,
161+
"Ip_profile": Ip_profile,
162+
"Phi": phi,
163+
"R_in": R_in,
164+
"R_out": R_out,
165+
"F": F,
166+
"int_dl_over_Bp": int_dl_over_Bp,
167+
"flux_surf_avg_1_over_R2": flux_surf_avg_1_over_R2,
168+
"flux_surf_avg_RBp": flux_surf_avg_RBp,
169+
"flux_surf_avg_R2Bp2": flux_surf_avg_R2Bp2,
170+
"flux_surf_avg_Bp2": flux_surf_avg_Bp2,
171+
"delta_upper_face": IMAS_data.profiles_1d.triangularity_upper,
172+
"delta_lower_face": IMAS_data.profiles_1d.triangularity_lower,
173+
"elongation": IMAS_data.profiles_1d.elongation,
174+
"vpr": vpr,
175+
"n_rho": n_rho,
176+
"hires_factor": hires_factor,
177+
"z_magnetic_axis": z_magnetic_axis,
178+
}
179+
180+
181+
def _load_imas_data(
182+
uri: str,
183+
ids_name: str,
184+
geometry_directory: str | None = None,
185+
) -> ids_toplevel.IDSToplevel:
186+
"""Loads a full IDS for a given uri or path_name and a given ids_name."""
187+
geometry_directory = geometry_loader.get_geometry_dir(geometry_directory)
188+
uri = os.path.join(geometry_directory, uri)
189+
with imas.DBEntry(uri=uri, mode="r") as db:
190+
ids = db.get(ids_name=ids_name)
191+
return ids

0 commit comments

Comments
 (0)