Multi-level modeling support from AMRs

[djinnome]

Ben Gyori has now support for expressions in distribution parameters so we can generate multi-level models. We need to properly sort the expressions so they are evaluated in order.

[djinnome]

Code to sample from parameter distributions should probably go here:

https://github.com/ciemss/pyciemss/blob/08cf9fc354e1de858cf0854aa5ca1bd31c56b2c6/pyciemss/mira_integration/compiled_dynamics.py#L84

[djinnome]

Hi @SamWitty

I have a question about how to sample from a PyroSample object.

In the _compile_param_values_mira function below, I have topologically sorted the parameters in mira, and now I am compiling each parameter in order:

https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/mira_integration/compiled_dynamics.py#L83-L110

In particular, on line 99, I pass the values of the parameters that have already been compiled to Pyro as free symbols that are used to parse the sympy expressions:

https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/mira_integration/compiled_dynamics.py#L99

However, to my surprise, the values dictionary that contains the parameter values has not been evaluated yet, so instead of the broadcast_tensors(*tensors) function receiving a list of tensors as its argument, it gets a list containing a PyroSample object instead:

pyciemss/compiled_dynamics.py:26: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../../../.pyenv/versions/miniconda3-3.11-24.1.2-0/envs/pyciemss/lib/python3.12/functools.py:909: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
pyciemss/mira_integration/compiled_dynamics.py:99: in _compile_param_values_mira
    param_value = mira_distribution_to_pyro(param_dist, free_symbols=values)
pyciemss/mira_integration/distributions.py:419: in mira_distribution_to_pyro
    k: safe_sympytorch_parse_expr(v, local_dict=free_symbols)
pyciemss/mira_integration/distributions.py:53: in safe_sympytorch_parse_expr
    return sympytorch.SymPyModule(expressions=[expr.args[0]])(**local_dict).squeeze()
../../../../.pyenv/versions/miniconda3-3.11-24.1.2-0/envs/pyciemss/lib/python3.12/site-packages/torch/nn/modules/module.py:1511: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
../../../../.pyenv/versions/miniconda3-3.11-24.1.2-0/envs/pyciemss/lib/python3.12/site-packages/torch/nn/modules/module.py:1520: in _call_impl
    return forward_call(*args, **kwargs)
../../../../.pyenv/versions/miniconda3-3.11-24.1.2-0/envs/pyciemss/lib/python3.12/site-packages/sympytorch/sympy_module.py:265: in forward
    out = torch.broadcast_tensors(*out)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

tensors = (PyroSample(prior=Beta()),)

    def broadcast_tensors(*tensors):
        r"""broadcast_tensors(*tensors) -> List of Tensors

        Broadcasts the given tensors according to :ref:`broadcasting-semantics`.

        Args:
            *tensors: any number of tensors of the same type

        .. warning::

            More than one element of a broadcasted tensor may refer to a single
            memory location. As a result, in-place operations (especially ones that
            are vectorized) may result in incorrect behavior. If you need to write
            to the tensors, please clone them first.

        Example::

            >>> x = torch.arange(3).view(1, 3)
            >>> y = torch.arange(2).view(2, 1)
            >>> a, b = torch.broadcast_tensors(x, y)
            >>> a.size()
            torch.Size([2, 3])
            >>> a
            tensor([[0, 1, 2],
                    [0, 1, 2]])
        """
        # This wrapper exists to support variadic args.
        if has_torch_function(tensors):
            return handle_torch_function(broadcast_tensors, tensors, *tensors)
>       return _VF.broadcast_tensors(tensors)  # type: ignore[attr-defined]
E       TypeError: expected Tensor as element 0 in argument 0, but got PyroSample

The reason for this is described in the Pyro documentation:

https://docs.pyro.ai/en/stable/nn.html#pyro.nn.module.PyroSample

assert isinstance(my_module, PyroModule)
my_module.x = PyroSample(Normal(0, 1))                    # independent
my_module.y = PyroSample(lambda self: Normal(self.x, 1))  # dependent

Note that my_module.y will not evaluate the lambda expression in the PyroSample object until getattr or setattr is called. Similarly, in the _compile_param_values_mira function, the pyro.distributions.Distribution object on line 103 is wrapped in a PyroSample object on line 104:

https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/mira_integration/compiled_dynamics.py#L103-L104

If the pyro.distributions.Distribution object contains unevaluated dependencies, such as the dependency of beta_mean on gamma_mean in the example below, then when should this dependency get evaluated?

beta_mean = Parameter(name='beta_mean',
                    distribution=Distribution(type="Beta1",
                    parameters={'alpha': sympy.Integer(10)*sympy.Symbol("gamma_mean"),
                                'beta': sympy.Integer(10)}))
gamma_mean = Parameter(name='gamma_mean',
                    distribution=Distribution(type="InverseGamma1",
                    parameters={'alpha': sympy.Integer(10),
                                'beta': sympy.Integer(10)}))

The unevaluated PyroSample objects are currently being held in the values dictionary, so when should I compile the values dictionary to a dictionary of tensors (let's call it compiled_values) so that I can compile the parameters that depend on their values?

Here are some options:

1. inside mira_distributions_to_pyro() where mira distribution parameters are passed to their corresponding pyro distribution function: https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/mira_integration/distributions.py#L418-L423
2. Inside _compile_param_values_mira() where distributions are wrapped in a PyroSample: https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/mira_integration/compiled_dynamics.py#L103-L104
3. Inside the pyciemss.compiled_dynamics.CompiledDynamics.__init__() method where samples are evaluated: https://github.com/ciemss/pyciemss/blob/e48ff2f88ff67f3c5e18b4916d3505d8f17f3bb4/pyciemss/compiled_dynamics.py#L36-L39

It seems that I would need to perform options 1-3 on gamma_mean parameter before I could perform options 1-3 on the beta_mean parameter

Any advice would be appreciated.