Radiation Sources

[allen-adastra]

Excited to get my hands on this!

I noticed radiation sources aren't modeled yet. It seems Brems + Synchrotron radiation should be pretty simple to add; perhaps just using the same formulas as implemented in cfspopcon

Thoughts on me opening a PR as a simple starter task?

[jcitrin]

Hi Allen! Indeed, Brems and Synchrotron radiation aren't in there yet, and they were short term goals. I'll add issues for the short-term roadmap for visibility. If you're interested it would be great if you open a PR with this. I'd also discuss with @araju around how to incorporate this in the config. We'd likely want the Stott formula to be one model option, keeping things open for other model options like the Wesson formula (which RAPTOR uses). The config about to be significantly refactored: @araju can provide further details.

[allen-adastra]

@araju would be good to hear more about the config refactor.

I've got a particular config system POPSIM side to build vectorized configs for (e.g. random walk generators, combinatorial cases, etc.); would be good to keep similarity where possible.

[allen-adastra]

One thing I've been wondering is what the motivation of using the DynamicConfigSlice pattern is over just doing a tree_map that resolves time-dependent quantities? Example:

def resolve_paths(tree: PyTree[typing.Union[ArrayLike, diffrax.AbstractPath]], t0: float, *args) -> PyTree[ArrayLike]:
    """Resolve any AbstractPath objects in a PyTree to their values at time t0.

    Args:
        tree (PyTree[Union[ArrayLike, AbstractPath]]): A PyTree of ArrayLike and AbstractPath objects.
        t0 (float): The time at which to evaluate the paths.

    Returns:
        PyTree[ArrayLike]: A PyTree of ArrayLike with all AbstractPath objects resolved to their values at time t0.
    """
    tree_resolved = jax.tree_map(
        lambda leaf: leaf.evaluate(t0, *args) if isinstance(leaf, diffrax.AbstractPath) else leaf,
        tree,
        is_leaf=lambda x: isinstance(x, diffrax.AbstractPath),
    )
    return tree_resolved

[araju]

@araju would be good to hear more about the config refactor.

I've got a particular config system POPSIM side to build vectorized configs for (e.g. random walk generators, combinatorial cases, etc.); would be good to keep similarity where possible.

tl;dr - we want the entire sim run to be configurable via "primitives" without having to write actual python code or deal with too many imports, etc. Current thinking is to use Fiddle for this Sim object.

Motivation: Right now, the simulation runs are configured via files like this where you have to define a get_sim() method which calls get_config() and get_geometry(). There are 2 issues with this:

1. The Config defined in there is a bit unweildy. Its structure is a holdover from PINT which TORAX is heavily based on. We need to have some more clear separation between runtime parameters and "static" config parameters. Also, the source params are a bit all over the place which need organizing.
2. We want to make things more user-friendly so that people don't need to, say, import a new stepper if they want to change from a linear solver to newton-raphson. For people less familiar with Python, we want to provide a Fiddle-based entrypoint where all they see is a json-like config where they can easily change things with primitives instead of instantiating actual objects.

This is a WIP, and actually Ian is helping on it now as well.

[araju]

One thing I've been wondering is what the motivation of using the DynamicConfigSlice pattern is over just doing a tree_map that resolves time-dependent quantities? Example:

def resolve_paths(tree: PyTree[typing.Union[ArrayLike, diffrax.AbstractPath]], t0: float, *args) -> PyTree[ArrayLike]:
    """Resolve any AbstractPath objects in a PyTree to their values at time t0.

    Args:
        tree (PyTree[Union[ArrayLike, AbstractPath]]): A PyTree of ArrayLike and AbstractPath objects.
        t0 (float): The time at which to evaluate the paths.

    Returns:
        PyTree[ArrayLike]: A PyTree of ArrayLike with all AbstractPath objects resolved to their values at time t0.
    """
    tree_resolved = jax.tree_map(
        lambda leaf: leaf.evaluate(t0, *args) if isinstance(leaf, diffrax.AbstractPath) else leaf,
        tree,
        is_leaf=lambda x: isinstance(x, diffrax.AbstractPath),
    )
    return tree_resolved

Great question. Basically, the input Config has 2 types of parameters: params that may be time-varying but really can simply change without recompilation, and params that must trigger a recompilation when they change. And it would be nice if the user-facing Config object did not have to be completely JAX-friendly.

In an effort to make it super clear the difference between these 2 types, we created the static and dynamic config slices. It also helps jit things while explicitly calling out the static args.

Now, that being said, the actual resolution to a specific value can be done the way you showed above. That's a good point. But the existence of the dynamic and static configs are more to help separate things for jitting.

Does that make sense?

[allen-adastra]

One thing I've been wondering is what the motivation of using the DynamicConfigSlice pattern is over just doing a tree_map that resolves time-dependent quantities? Example:

def resolve_paths(tree: PyTree[typing.Union[ArrayLike, diffrax.AbstractPath]], t0: float, *args) -> PyTree[ArrayLike]:
    """Resolve any AbstractPath objects in a PyTree to their values at time t0.

    Args:
        tree (PyTree[Union[ArrayLike, AbstractPath]]): A PyTree of ArrayLike and AbstractPath objects.
        t0 (float): The time at which to evaluate the paths.

    Returns:
        PyTree[ArrayLike]: A PyTree of ArrayLike with all AbstractPath objects resolved to their values at time t0.
    """
    tree_resolved = jax.tree_map(
        lambda leaf: leaf.evaluate(t0, *args) if isinstance(leaf, diffrax.AbstractPath) else leaf,
        tree,
        is_leaf=lambda x: isinstance(x, diffrax.AbstractPath),
    )
    return tree_resolved

Great question. Basically, the input Config has 2 types of parameters: params that may be time-varying but really can simply change without recompilation, and params that must trigger a recompilation when they change. And it would be nice if the user-facing Config object did not have to be completely JAX-friendly.

In an effort to make it super clear the difference between these 2 types, we created the static and dynamic config slices. It also helps jit things while explicitly calling out the static args.

Now, that being said, the actual resolution to a specific value can be done the way you showed above. That's a good point. But the existence of the dynamic and static configs are more to help separate things for jitting.

Does that make sense?

Gotcha; cool thanks! Yeah I did something similar with having a class Config for static "compile-time" configuration and class Params for (possibly time-dependent) "run-time" configuration. Although I'm not super confident that the way I set things up avoids all the foot-guns I would like to. I got your email, would be good to exchange notes!

[jcitrin]

Bremsstrahlung added in #176 . Will keep this issue open for cyclotron and line radiation source discussions.