Coverage for src/hallmd/models/thruster.py: 79%
183 statements
« prev ^ index » next coverage.py v7.6.10, created at 2025-02-20 20:43 +0000
1"""Module for Hall thruster models.
3!!! Note
4 Only current implementation is for the 1d fluid [Hallthruster.jl code](https://github.com/UM-PEPL/HallThruster.jl).
5 Other thruster codes can be implemented similarly here.
7Includes:
9- `run_hallthruster_jl` - General wrapper to run HallThruster.jl for a single set of inputs
10- `hallthruster_jl()` - PEM wrapper to run HallThruster.jl for a set of PEM inputs
11- `get_jl_env` - Get the path of the julia environment created for HallThruster.jl for a specific git ref
12- `PEM_TO_JULIA` - Mapping of PEM variable names to a path in the HallThruster.jl input/output structure (defaults)
13"""
15import copy
16import json
17import os
18import random
19import string
20import subprocess
21import tempfile
22import time
23import warnings
24from importlib import resources
25from pathlib import Path
26from typing import Callable, Literal
28import numpy as np
29from amisc.typing import Dataset
31from hallmd.utils import AVOGADRO_CONSTANT, FUNDAMENTAL_CHARGE, MOLECULAR_WEIGHTS, load_device
33__all__ = ["run_hallthruster_jl", "hallthruster_jl", "get_jl_env", "PEM_TO_JULIA"]
35HALLTHRUSTER_VERSION_DEFAULT = "0.18.3"
37# Maps PEM variable names to a path in the HallThruster.jl input/output structure (default values here)
38with open(resources.files("hallmd.models") / "pem_to_julia.json", "r") as fd:
39 PEM_TO_JULIA = json.load(fd)
42def get_jl_env(git_ref: str) -> Path:
43 """Get the path of the julia environment created for HallThruster.jl for a specific git ref.
45 :param git_ref: The git ref (i.e. commit hash, version tag, branch, etc.) of HallThruster.jl to use.
46 """
47 global_env_dir = Path("~/.julia/environments/").expanduser()
48 env_path = global_env_dir / f"hallthruster_{git_ref}"
49 return env_path
52def _convert_to_julia(pem_data: dict, julia_data: dict, pem_to_julia: dict):
53 """Replace all values in the mutable `julia_data` dict with corresponding values in `pem_data`, using the
54 conversion map provided in `pem_to_julia`. Will "blaze" a path into `julia_data` dict if it does not exist.
56 :param pem_data: thruster inputs of the form `{'pem_variable_name': value}`
57 :param julia_data: a `HallThruster.jl` data structure of the form `{'config': {...}, 'simulation': {...}, etc.}`
58 :param pem_to_julia: a conversion map from a pem variable name to a path in `julia_data`. Use strings for dict keys
59 and integers for list indices. For example, `{'P_b': ['config', 'background_pressure', 2]}`
60 will set `julia_data['config']['background_pressure'][2] = pem_data['P_b']`.
61 """
62 for pem_key, value in pem_data.items():
63 if pem_key not in pem_to_julia:
64 raise KeyError(f"Cannot convert PEM data variable {pem_key} since it is not in the provided conversion map")
66 # Blaze a trail into the julia data structure (set dict defaults for str keys and list defaults for int keys)
67 julia_path = pem_to_julia.get(pem_key)
68 pointer = julia_data
69 for i, key in enumerate(julia_path[:-1]):
70 if isinstance(pointer, dict) and not pointer.get(key):
71 pointer.setdefault(key, {} if isinstance(julia_path[i + 1], str) else [])
72 if isinstance(pointer, list) and len(pointer) <= key:
73 pointer.extend(
74 [{} if isinstance(julia_path[i + 1], str) else [] for _ in range(key - len(pointer) + 1)]
75 )
76 pointer = pointer[key]
77 pointer[julia_path[-1]] = value
80def _convert_to_pem(julia_data: dict, pem_to_julia: dict):
81 """Return a `dict` of PEM outputs from a `HallThruster.jl` output structure, using the conversion map provided."""
82 pem_data = {}
83 for pem_key, julia_path in pem_to_julia.items():
84 pointer = julia_data
85 if julia_path[0] == "output":
86 found_value = True
87 for key in julia_path:
88 try:
89 pointer = pointer[key]
90 except (KeyError, IndexError):
91 found_value = False
92 break
94 if found_value:
95 pem_data[pem_key] = pointer
96 return pem_data
99def _default_model_fidelity(model_fidelity: tuple, json_config: dict, cfl: float = 0.2) -> dict:
100 """Built-in (default) method to convert model fidelity tuple to `ncells` and `ncharge` via:
102 ```python
103 ncells = 50 * (model_fidelity[0] + 2)
104 ncharge = model_fidelity[1] + 1
105 ```
107 Also adjusts the time step `dt` to maintain the CFL condition (based on grid spacing and ion velocity).
109 :param model_fidelity: tuple of integers that determine the number of cells and the number of charge states to use
110 :param json_config: the current set of configurations for HallThruster.jl
111 :param cfl: a conservative estimate for CFL condition to determine time step
112 :returns: a dictionary of simulation parameters that can be converted to Julia via the `pem_to_julia` mapping,
113 namely `{'num_cells': int, 'ncharge': int, 'dt': float}`
114 """
115 if model_fidelity == ():
116 model_fidelity = (2, 2) # default to high-fidelity model
118 num_cells = 50 * (model_fidelity[0] + 2)
119 ncharge = model_fidelity[1] + 1
121 # Estimate conservative time step based on CFL and grid spacing
122 config = json_config.get("config", {})
123 domain = config.get("domain", [0, 0.08])
124 anode_pot = config.get("discharge_voltage", 300)
125 cathode_pot = config.get("cathode_coupling_voltage", 0)
126 propellant = config.get("propellant", "Xenon")
128 if propellant not in MOLECULAR_WEIGHTS:
129 warnings.warn(
130 f"Could not find propellant {propellant} in `hallmd.utils`. "
131 f"Will default to Xenon for estimating uniform time step from CFL..."
132 )
133 propellant = "Xenon"
135 mi = MOLECULAR_WEIGHTS[propellant] / AVOGADRO_CONSTANT / 1000 # kg
136 dx = float(domain[1]) / (num_cells + 1)
137 u = np.sqrt(2 * ncharge * FUNDAMENTAL_CHARGE * (anode_pot - cathode_pot) / mi)
138 dt_s = cfl * dx / u
140 return {"num_cells": num_cells, "ncharge": ncharge, "dt": float(dt_s)}
143def _format_hallthruster_jl_input(
144 thruster_inputs: dict,
145 thruster: dict | str = "SPT-100",
146 config: dict = None,
147 simulation: dict = None,
148 postprocess: dict = None,
149 model_fidelity: tuple = (2, 2),
150 output_path: str | Path = None,
151 pem_to_julia: dict = "default",
152 fidelity_function: Callable = "default",
153) -> dict:
154 """Helper function to format PEM inputs for `Hallthruster.jl` as a `dict` writeable to `json`. See the call
155 signature of `hallthruster_jl` for more details on arguments.
157 The return `dict` has the format:
158 ```json
159 {
160 "config": {
161 "thruster": {...},
162 "discharge_voltage": 300,
163 "domain": [0, 0.08],
164 "anode_mass_flow_rate": 1e-6,
165 "background_pressure": 1e-5,
166 etc.
167 },
168 "simulation": {
169 "ncells": 202,
170 "dt": 8.4e-9
171 "duration": 1e-3,
172 etc.
173 }
174 "postprocess": {...}
175 }
176 ```
178 :returns: a json `dict` in the format that `HallThruster.run_simulation()` expects to be called
179 """
180 if fidelity_function is None or fidelity_function == "default":
181 fidelity_function = _default_model_fidelity
183 json_config = {
184 "config": {} if config is None else copy.deepcopy(config),
185 "simulation": {} if simulation is None else copy.deepcopy(simulation),
186 "postprocess": {} if postprocess is None else copy.deepcopy(postprocess),
187 }
189 # Necessary to load thruster specs separately from the config (to protect sensitive data)
190 if isinstance(thruster, str):
191 thruster = load_device(thruster)
192 if thruster is not None:
193 json_config["config"]["thruster"] = thruster # override
195 # Make sure we request time-averaged
196 duration = json_config["simulation"].get("duration", 1e-3)
197 avg_start_time = json_config["postprocess"].get("average_start_time", 0.5 * duration)
198 json_config["postprocess"]["average_start_time"] = avg_start_time
200 # Update/override config with PEM thruster inputs (modify in place)
201 _convert_to_julia(thruster_inputs, json_config, pem_to_julia)
203 # Override model fidelity quantities
204 if model_fidelity is not None:
205 fidelity_overrides = fidelity_function(model_fidelity, json_config)
206 _convert_to_julia(fidelity_overrides, json_config, pem_to_julia)
208 if output_path is not None:
209 fname = "hallthruster_jl"
210 if name := json_config["config"].get("thruster", {}).get("name"):
211 fname += f"_{name}"
212 if vd := json_config["config"].get("discharge_voltage"):
213 fname += f"_{round(vd)}V"
214 if mdot := json_config["config"].get("anode_mass_flow_rate"):
215 fname += f"_{mdot:.1e}kg_s"
217 fname += "_" + "".join(random.choices(string.ascii_uppercase + string.digits, k=4)) + ".json"
218 output_file = str((Path(output_path) / fname).resolve())
219 json_config["postprocess"]["output_file"] = output_file
221 # Handle special conversions for anomalous transport models
222 if anom_model := json_config["config"].get("anom_model"):
223 if anom_model.get("type") in ["LogisticPressureShift", "SimpleLogisticShift"]:
224 anom_model = anom_model.get("model", {})
226 match anom_model.get("type", "TwoZoneBohm"):
227 case "TwoZoneBohm":
228 if thruster_inputs.get("a_2") is not None: # Only when the PEM a_2 is used
229 anom_model["c2"] = anom_model["c2"] * anom_model.get("c1", 0.00625)
230 case "GaussianBohm":
231 if thruster_inputs.get("anom_max") is not None: # Only when the PEM anom_max is used
232 anom_model["hall_max"] = anom_model["hall_max"] * anom_model.get("hall_min", 0.00625)
234 return json_config
237def run_hallthruster_jl(
238 json_input: dict | str | Path, jl_env: str | Path = None, jl_script: str | Path = None, **kwargs
239) -> dict:
240 """Python wrapper for `HallThruster.run_simulation(json_input)` in Julia.
242 :param json_input: either a dictionary containing `config`, `simulation`, and `postprocess` options for
243 HallThruster.jl, or a string/Path containing a path to a JSON file with those inputs.
244 :param jl_env: The julia environment containing HallThruster.jl. Defaults to global Julia environment.
245 :param jl_script: path to a custom Julia script to run. The script should accept the input json file path as
246 a command line argument. Defaults to just calling `HallThruster.run_simulation(input_file)`.
247 :param kwargs: additional keyword arguments to pass to `subprocess.run` when calling the Julia script.
249 :returns: `dict` of `Hallthruster.jl` outputs. The specific outputs depend on the settings
250 provided in the `postprocess` dict in the input. If `postprocess['output_file']` is present,
251 this function will also write the requested outputs and restart information to that file.
252 """
253 # Read JSON input from file if path provided
254 if isinstance(json_input, str | Path):
255 with open(json_input, "r") as fp:
256 json_input = json.load(fp)
258 tempfile_args = dict(suffix=".json", prefix="hallthruster_jl_", mode="w", delete=False, encoding="utf-8")
260 # Get output file path. If one not provided, create a temporary
261 temp_out = False
262 if "output_file" in json_input.get("postprocess", {}):
263 output_file = Path(json_input["postprocess"].get("output_file"))
264 elif "output_file" in json_input.get("input", {}).get("postprocess", {}):
265 output_file = Path(json_input["input"]["postprocess"].get("output_file"))
266 else:
267 temp_out = True
268 fd_out = tempfile.NamedTemporaryFile(**tempfile_args)
269 output_file = Path(fd_out.name)
270 fd_out.close()
272 if json_input.get("input"):
273 json_input["input"].setdefault("postprocess", {})
274 json_input["input"]["postprocess"]["output_file"] = str(output_file.resolve())
275 else:
276 json_input.setdefault("postprocess", {})
277 json_input["postprocess"]["output_file"] = str(output_file.resolve())
279 # Dump input to temporary file
280 fd = tempfile.NamedTemporaryFile(**tempfile_args)
281 input_file = fd.name
282 json.dump(json_input, fd, ensure_ascii=False, indent=4)
283 fd.close()
285 # Run HallThruster.jl on input file
286 if jl_script is None:
287 cmd = [
288 "julia",
289 "--startup-file=no",
290 "-e",
291 f'using HallThruster; HallThruster.run_simulation(raw"{input_file}")',
292 ]
293 else:
294 cmd = ["julia", "--startup-file=no", "--", str(Path(jl_script).resolve()), input_file]
296 if jl_env is not None:
297 if Path(jl_env).exists():
298 cmd.insert(1, f"--project={Path(jl_env).resolve()}")
299 else:
300 raise ValueError(
301 f"Could not find Julia environment {jl_env}. Please create it first. "
302 f"See https://github.com/JANUS-Institute/HallThrusterPEM/blob/main/scripts/install_hallthruster.py"
303 )
305 try:
306 subprocess.run(cmd, **kwargs)
307 finally:
308 # Delete temporary input file
309 os.unlink(input_file)
311 # Load output data
312 with open(output_file, "r") as fp:
313 output_data = json.load(fp)
315 if temp_out:
316 os.unlink(output_file)
317 if d := output_data.get("postprocess"):
318 if "output_file" in d:
319 del d["output_file"]
320 if d := output_data.get("input"):
321 if d2 := d.get("postprocess"):
322 if "output_file" in d2:
323 del d2["output_file"]
325 return output_data
328def hallthruster_jl(
329 thruster_inputs: Dataset = None,
330 thruster: Literal["SPT-100"] | str | dict = "SPT-100",
331 config: dict = None,
332 simulation: dict = None,
333 postprocess: dict = None,
334 model_fidelity: tuple = (2, 2),
335 output_path: str | Path = None,
336 version: str = HALLTHRUSTER_VERSION_DEFAULT,
337 pem_to_julia: dict = "default",
338 fidelity_function: Callable[[tuple[int, ...]], dict] = "default",
339 julia_script: str | Path = None,
340 run_kwargs: dict = "default",
341 shock_threshold: float = None,
342) -> Dataset:
343 """Run a single `HallThruster.jl` simulation for a given set of inputs. This function will write a temporary
344 input file to disk, call `HallThruster.run_simulation()` in Julia, and read the output file back into Python. Will
345 return time-averaged performance metrics and ion velocity for use with the PEM.
347 Note that the specific inputs and outputs described here can be configured using the `pem_to_julia` dict.
349 !!! Warning "Required configuration"
350 You must specify a thruster, a domain, a mass flow rate, and a discharge voltage to run the simulation. The
351 thruster must be defined in the `hallmd.devices` directory or as a dictionary with the required fields.
352 The mass flow rate and discharge voltage are specified in `thruster_inputs` as `mdot_a` (kg/s) and
353 `V_a` (V), respectively. The domain is specified as a list `[left_bound, right_bound]` in the
354 `config` dictionary. See the
355 [HallThruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/reference/config/) for more details.
357 :param thruster_inputs: named key-value pairs of thruster inputs: `P_b`, `V_a`, `mdot_a`, `T_e`, `u_n`, `l_t`,
358 `a_1`, `a_2`, `delta_z`, `z0`, `p0`, and `V_cc` for background pressure (Torr), anode
359 voltage, anode mass flow rate (kg/s), electron temperature (eV), neutral velocity (m/s),
360 transition length (m), anomalous transport coefficients, and cathode coupling voltage. Will
361 override the corresponding values in `config` if provided.
362 :param thruster: the name of the thruster to simulate (must be importable from `hallmd.devices`, see
363 [`load_device`][hallmd.utils.load_device]), or a dictionary that provides geometry and
364 magnetic field information of the thruster to simulate; see the
365 [Hallthruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/tutorials/simulation//run/).
366 Will override `thruster` in `config` if provided. If None, will defer to `config`.
367 Defaults to the SPT-100.
368 :param config: dictionary of configs for `HallThruster.jl`, see the
369 [Hallthruster.jl docs](https://um-pepl.github.io/HallThruster.jl/dev/reference/config/) for
370 options and formatting.
371 :param simulation: dictionary of simulation parameters for `HallThruster.jl`
372 :param postprocess: dictionary of post-processing parameters for `Hallthruster.jl`
373 :param model_fidelity: tuple of integers that determine the number of cells and the number of charge states to use
374 via `ncells = model_fidelity[0] * 50 + 100` and `ncharge = model_fidelity[1] + 1`.
375 Will override `ncells` and `ncharge` in `simulation` and `config` if provided.
376 :param output_path: base path to save output files, will write to current directory if not specified
377 :param version: version of HallThruster.jl to use; will
378 search for a global `hallthruster_{version}` environment in the `~/.julia/environments/` directory.
379 Can also specify a specific git ref (i.e. branch, commit hash, etc.) to use from GitHub. If the
380 `hallthruster_{version}` environment does not exist, an error will be raised -- you should create
381 this environment first before using it.
382 :param pem_to_julia: a `dict` mapping of PEM shorthand variable names to a list of keys that maps into the
383 `HallThruster.jl` input/output data structure. Defaults to the provided PEM_TO_JULIA dict
384 defined in [`hallmd.models.thruster`][hallmd.models.thruster]. For example,
385 `{'P_b': ['config', 'background_pressure']}` will set `config['background_pressure'] = P_b`.
386 If specified, will override and extend the default mapping.
387 :param fidelity_function: a callable that takes a tuple of integers and returns a dictionary of simulation
388 parameters. Defaults to `_default_model_fidelity` which sets `ncells` and `ncharge` based
389 on the input tuple. The returned simulation parameters must be convertable to Julia via
390 the `pem_to_julia` mapping. The callable should also take in the current json config dict.
391 :param julia_script: path to a custom Julia script to run. The script should accept the input json file path as
392 a command line argument. Defaults to just calling `HallThruster.run_simulation(input_file)`.
393 :param run_kwargs: additional keyword arguments to pass to `subprocess.run` when calling the Julia script.
394 Defaults to `check=True`.
395 :param shock_threshold: if provided, an error will be raised if the ion velocity reaches a maximum before this
396 threshold axial location (in m) - used to detect and filter unwanted "shock-like" behavior,
397 for example by providing a threshold of half the domain length. If not provided, then no
398 filtering is performed (default).
399 :returns: `dict` of `Hallthruster.jl` outputs: `I_B0`, `I_d`, `T`, `eta_c`, `eta_m`, `eta_v`, and `u_ion` for ion
400 beam current (A), discharge current (A), thrust (N), current efficiency, mass efficiency, voltage
401 efficiency, and singly-charged ion velocity profile (m/s), all time-averaged.
402 """
403 if pem_to_julia is None or pem_to_julia == "default":
404 pem_to_julia = copy.deepcopy(PEM_TO_JULIA)
405 else:
406 tmp = copy.deepcopy(PEM_TO_JULIA)
407 tmp.update(pem_to_julia)
408 pem_to_julia = tmp
410 thruster_inputs = thruster_inputs or {}
412 # Format PEM inputs for HallThruster.jl
413 json_data = _format_hallthruster_jl_input(
414 thruster_inputs,
415 thruster=thruster,
416 config=config,
417 simulation=simulation,
418 postprocess=postprocess,
419 model_fidelity=model_fidelity,
420 output_path=output_path,
421 pem_to_julia=pem_to_julia,
422 fidelity_function=fidelity_function,
423 )
424 # Get julia environment
425 jl_environment = get_jl_env(version) if version is not None else None
427 if run_kwargs is None:
428 run_kwargs = {}
429 elif run_kwargs == "default":
430 run_kwargs = {"check": True}
432 # Run Julia
433 t1 = time.time()
434 sim_results = run_hallthruster_jl(json_data, jl_env=jl_environment, jl_script=julia_script, **run_kwargs)
435 t2 = time.time()
437 # Format QOIs for PEM
438 thruster_outputs = _convert_to_pem(sim_results, pem_to_julia)
440 # Raise an exception if thrust or beam current are negative (non-physical cases)
441 thrust = thruster_outputs.get("T", 0)
442 beam_current = thruster_outputs.get("I_B0", 0)
443 if thrust < 0 or beam_current < 0:
444 raise ValueError(f"Exception due to non-physical case: thrust={thrust} N, beam current={beam_current} A")
446 # Raise an exception for ion velocity that exhibits "shock" behavior
447 if shock_threshold is not None:
448 z_coords = thruster_outputs.get("u_ion_coords")
449 ion_velocity = thruster_outputs.get("u_ion")
450 if z_coords is not None and ion_velocity is not None:
451 if (z_max := z_coords[np.argmax(ion_velocity)]) < shock_threshold:
452 raise ValueError(f"Exception due to shock-like behavior: max ion velocity occurs at z={z_max:.3f} m")
454 thruster_outputs["model_cost"] = t2 - t1 # seconds
456 if output_path is not None:
457 output_file = Path(json_data["postprocess"].get("output_file"))
458 thruster_outputs["output_path"] = output_file.relative_to(Path(output_path).resolve()).as_posix()
460 return thruster_outputs