"""Allows initializing GRILLIX data from filepath."""
from pathlib import Path
from numbers import Integral
from typing import Optional, Sequence, Union
from warnings import warn
import torx
from torx.units import Quantity
from torx.fileio import filepath_resolver, read_fortran_namelist, search_paths
from torx.equilibrium import NumericalEquilibrium, initialize_equi_from_params
from torx.units import Normalization
from torx.decorators import autodoc_function
def _plane_grid_directory(filepath: Path) -> Optional[Path]:
"""
Find the directory holding the multigrid file of every plane.
The same directories are searched as by filepath_resolver, so a run
which holds its grids next to the output or in its trunk is found.
Parameters
----------
filepath : Path
Path of the GRILLIX output directory.
Returns
-------
Path or None
Directory holding the files, None for a run whose planes share a
single grid.
"""
for search_path in search_paths:
directory = filepath / search_path
if directory.exists() and any(directory.glob("multigrids_plane*.nc")):
return directory
return None
[docs]
@autodoc_function
def initialize_grillix_from_filepath(
filepath: Path,
planes: Union[int, str, Sequence[int]] = 0
):
"""
Initialize GRILLIX data from filepath.
Parameters
----------
filepath : Path
Path of the GRILLIX output directory.
planes : int, str or Sequence[int]
Plane of a run holding a grid per plane to return the 2D grid of.
Give a list of indices for a 3D grid holding those planes, or 'all'
for a 3D grid holding every plane. A run whose planes share a
single grid returns that grid for any index.
Returns
-------
tuple
The grid, the equilibrium, the parameters and the normalization.
Asking for more than one plane returns the canonical and the
staggered 3D grid, so five values instead of four.
Raises
------
ValueError
If planes is a string other than 'all', or if more than one
plane is asked for of a run whose planes share a single grid.
"""
filepath = Path(filepath)
if isinstance(planes, str):
if planes.lower() != "all":
raise ValueError(
f"Got planes='{planes}', expected the index of a plane, "
"a list of indices or 'all' for the whole 3D grid."
)
wanted_planes = None
elif isinstance(planes, Integral):
wanted_planes = int(planes)
else:
wanted_planes = [int(index) for index in planes]
try:
params_filepath = filepath_resolver(filepath, "params_static_data.nml")
except FileNotFoundError:
line1 = f"'params_static_data.nml' not found in {filepath}"
line2 = " -> falling back to legacy 'params.in'"
print(f"{line1}\n{line2}")
params_filepath = filepath_resolver(filepath, "params.in")
params = read_fortran_namelist(params_filepath)
params_braginskii = read_fortran_namelist(
filepath_resolver(filepath, "params_braginskii.in")
)
# Combine the two parameter dictionaries
params = {**params, **params_braginskii}
equi = initialize_equi_from_params(filepath, params)
try:
norm = Normalization.initialize_from_normalization_file(
filepath_resolver(filepath, "physical_parameters.nml")
)
equi.B0 = norm.B0
R0 = norm.R0
except FileNotFoundError:
warn(
"physical_parameters.nml not found. Normalization will not be "
f"usable. Using 'axis_Btor' as B0 from {equi.filepath}."
)
norm = Normalization(dict())
equi.B0 = equi.axis_Btor
R0 = Quantity(1, "m")
if (
isinstance(equi, NumericalEquilibrium)
and params["equi_numerical_params"]["flip_z"]
):
equi.flip_Z()
plane_directory = _plane_grid_directory(filepath)
if plane_directory is None:
# Returning the single grid here would give four values instead of
# the five of a 3D grid, which the caller cannot unpack.
if not isinstance(wanted_planes, int):
raise ValueError(
"More than one plane was asked for, but the planes of this "
"run share a single grid, so there is no 3D grid to load."
)
if wanted_planes != 0:
warn(
f"Plane {planes} was asked for, but the planes of this run "
"share a single grid. That grid is returned."
)
multigrid_file = filepath_resolver(filepath, "multigrid.nc")
multigrid = torx.grid.Multigrid2D(
multigrid_file,
multigrid_group = "",
load_all_levels = False,
R0 = R0
)
grid = multigrid.get_grid(1)
return grid, equi, params, norm
# A single plane is taken out of the whole grid below, so only a list
# of planes limits which files are read.
wanted = wanted_planes if isinstance(wanted_planes, list) else None
grid_cano = torx.grid.Grid3D.from_multigrid_files(plane_directory,
planes = wanted)
grid_stag = torx.grid.Grid3D.from_multigrid_files(plane_directory,
planes = wanted,
staggered = True)
grid_cano.set_R0(R0)
grid_stag.set_R0(R0)
if isinstance(wanted_planes, int):
return grid_cano.isel_phi(wanted_planes), equi, params, norm
return grid_cano, grid_stag, equi, params, norm
