google-deepmind / torax

TORAX: Tokamak transport simulation in JAX
https://torax.readthedocs.io
Other
297 stars 26 forks source link

Make jtot_hires optional #220

Closed theo-brown closed 1 week ago

theo-brown commented 1 week ago

Fixes issue #219.

Note that although the existing tests pass, I have not added any new tests to check this functionality. However, this should not be a problem, as the bit I removed was permanently unused anyway.

theo-brown commented 1 week ago

One thing I wasn't certain of was whether setting a frozen Chex dataclass field to Optional was permitted.

jcitrin commented 1 week ago

One thing I wasn't certain of was whether setting a frozen Chex dataclass field to Optional was permitted.

Shouldn't be an issue. See other parts of TORAX where we do this, e.g. https://github.com/google-deepmind/torax/blob/b13133f9a3bd4a4a8fd3ab4ba71949244453a67c/torax/fvm/cell_variable.py#L67