Merge pull request #386 from mj-will/support-glasflow-flows

Add experimental support for glasflow flows

Author: mj-will

codecov.yml

@@ -4,3 +4,5 @@ coverage:
       default:
         target: 90%
         threshold: 1%
+  ignore:
+    - "nessai/experimental"

nessai/experimental/__init__.py

@@ -0,0 +1,83 @@
+from functools import partial
+from glasflow.flows import CouplingNSF, RealNVP
+
+from ...flows.base import BaseFlow
+
+known_flows = {
+    "nsf": CouplingNSF,
+    "realnvp": RealNVP,
+}
+
+
+class GlasflowWrapper(BaseFlow):
+    """Wrapper for glasflow flow classes"""
+
+    def __init__(
+        self,
+        FlowClass,
+        n_inputs,
+        n_neurons,
+        n_blocks,
+        n_layers,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        n_conditional_inputs = kwargs.pop("context_features", None)
+        self._flow = FlowClass(
+            n_inputs=n_inputs,
+            n_transforms=n_blocks,
+            n_blocks_per_transform=n_layers,
+            n_neurons=n_neurons,
+            n_conditional_inputs=n_conditional_inputs,
+            **kwargs,
+        )
+
+    def forward(self, x, context=None):
+        return self._flow.forward(x, conditional=context)
+
+    def inverse(self, z, context=None):
+        return self._flow.inverse(z, conditional=context)
+
+    def log_prob(self, x, context=None):
+        return self._flow.log_prob(x, conditional=context)
+
+    def sample(self, n, context=None):
+        return self._flow.sample(n, conditional=context)
+
+    def forward_and_log_prob(self, x, context=None):
+        return self._flow.forward_and_log_prob(x, conditional=context)
+
+    def sample_and_log_prob(self, n, context=None):
+        return self._flow.sample_and_log_prob(n, conditional=context)
+
+    def sample_latent_distribution(self, n, context=None):
+        if context is not None:
+            raise ValueError
+        return self._flow.distribution.sample(n)
+
+    def base_distribution_log_prob(self, z, context=None):
+        if context is not None:
+            raise ValueError("Context must be None")
+        return self._flow.base_distribution_log_prob(z)
+
+    def freeze_transform(self):
+        self._flow._transform.requires_grad_(False)
+
+    def unfreeze_transform(self):
+        self._flow._transform.requires_grad_(True)
+
+
+def get_glasflow_class(name):
+    """Get the class for a glasflow flow.
+
+    Note: the name must start with :code:`glasflow.`
+    """
+    name = name.lower()
+    if "glasflow" not in name:
+        raise ValueError("'glasflow' missing from name")
+    short_name = name.replace("glasflow-", "")
+    if short_name not in known_flows:
+        raise ValueError(f"{name} is not a known glasflow flow")
+    FlowClass = known_flows.get(short_name)
+    return partial(GlasflowWrapper, FlowClass)

nessai/experimental/flows/__init__.py

@@ -165,23 +165,45 @@ def get_n_neurons(
     return n
 
 
-def configure_model(config):
-    """
-    Setup the flow form a configuration dictionary.
-    """
+def get_native_flow_class(name):
+    """Get a natively implemented flow class."""
+    name = name.lower()
     from .realnvp import RealNVP
     from .maf import MaskedAutoregressiveFlow
     from .nsf import NeuralSplineFlow
-    from ..flowmodel import config as fmconfig
 
-    kwargs = {}
     flows = {
         "realnvp": RealNVP,
         "maf": MaskedAutoregressiveFlow,
         "frealnvp": RealNVP,
         "spline": NeuralSplineFlow,
         "nsf": NeuralSplineFlow,
     }
+    if name not in flows:
+        raise ValueError(f"Unknown flow: {name}")
+    return flows.get(name)
+
+
+def get_flow_class(name: str):
+    """Get the class to use for the normalizing flow from a string."""
+    name = name.lower()
+    if "glasflow" in name:
+        from ..experimental.flows.glasflow import get_glasflow_class
+
+        logger.warning("Using experimental glasflow flow!")
+        FlowClass = get_glasflow_class(name)
+    else:
+        FlowClass = get_native_flow_class(name)
+    return FlowClass
+
+
+def configure_model(config):
+    """
+    Setup the flow form a configuration dictionary.
+    """
+    from ..flowmodel import config as fmconfig
+
+    kwargs = {}
     activations = {"relu": F.relu, "tanh": F.tanh, "swish": silu, "silu": silu}
 
     config = config.copy()
@@ -218,39 +240,21 @@ def configure_model(config):
     if distribution:
         kwargs["distribution"] = distribution
 
-    fc = config.get("flow", None)
-    ftype = config.get("ftype", None)
-    if fc is not None:
-        model = fc(
-            config["n_inputs"],
-            config["n_neurons"],
-            config["n_blocks"],
-            config["n_layers"],
-            **kwargs,
-        )
-    elif ftype is not None:
-        if ftype.lower() not in flows:
-            raise RuntimeError(
-                f"Unknown flow type: {ftype}. Choose from:" f"{flows.keys()}"
-            )
-        if (
-            ("mask" in kwargs and kwargs["mask"] is not None)
-            or ("net" in kwargs and kwargs["net"] is not None)
-        ) and ftype.lower() not in ["realnvp", "frealnvp"]:
-            raise RuntimeError(
-                "Custom masks and networks are only " "supported for RealNVP"
-            )
-
-        model = flows[ftype.lower()](
-            config["n_inputs"],
-            config["n_neurons"],
-            config["n_blocks"],
-            config["n_layers"],
-            **kwargs,
-        )
-    else:
+    FlowClass = config.get("flow")
+    ftype = config.get("ftype")
+    if FlowClass is None and ftype is None:
         raise RuntimeError("Must specify either 'flow' or 'ftype'.")
 
+    if FlowClass is None:
+        FlowClass = get_flow_class(ftype)
+    model = FlowClass(
+        config["n_inputs"],
+        config["n_neurons"],
+        config["n_blocks"],
+        config["n_layers"],
+        **kwargs,
+    )
+
     device = torch.device(config.get("device_tag", "cpu"))
     if device != "cpu":
         try:

nessai/experimental/flows/glasflow.py

@@ -0,0 +1,49 @@
+from nessai.flowmodel import FlowModel
+from nessai.experimental.flows.glasflow import (
+    GlasflowWrapper,
+    get_glasflow_class,
+    known_flows,
+)
+import numpy as np
+import pytest
+
+
+@pytest.mark.parametrize("name", known_flows.keys())
+def test_get_glasflow_class(name):
+    FlowClass = get_glasflow_class(f"glasflow-{name}")
+    FlowClass(n_inputs=2, n_neurons=4, n_blocks=2, n_layers=1)
+
+
+def test_get_glasflow_class_missing_prefix():
+    with pytest.raises(ValueError, match=r"'glasflow' missing from name"):
+        get_glasflow_class("realnvp")
+
+
+def test_get_glasflow_class_invalid_flow():
+    with pytest.raises(
+        ValueError, match=r"invalid is not a known glasflow flow"
+    ):
+        get_glasflow_class("glasflow.invalid")
+
+
+@pytest.mark.integration_test
+def test_glasflow_integration(tmp_path):
+
+    from glasflow.flows import RealNVP
+
+    config = dict(
+        model_config=dict(
+            ftype="glasflow-realnvp",
+            n_inputs=2,
+            kwargs=None,
+        )
+    )
+
+    flowmodel = FlowModel(config=config, output=tmp_path / "test")
+
+    flowmodel.initialise()
+
+    assert isinstance(flowmodel.model, GlasflowWrapper)
+    assert isinstance(flowmodel.model._flow, RealNVP)
+
+    flowmodel.train(np.random.randn(100, 2))

nessai/flows/utils.py

@@ -12,11 +12,18 @@
     create_linear_transform,
     create_pre_transform,
     get_base_distribution,
+    get_flow_class,
     get_n_neurons,
+    get_native_flow_class,
     silu,
     reset_weights,
     reset_permutations,
 )
+from nessai.flows import (
+    RealNVP,
+    NeuralSplineFlow,
+    MaskedAutoregressiveFlow,
+)
 
 
 @pytest.fixture
@@ -196,6 +203,47 @@ def test_reset_permutation_lu():
     lu._initialize.assert_called_once_with(identity_init=True)
 
 
+@pytest.mark.parametrize(
+    "name, expected_class",
+    [
+        ("realnvp", RealNVP),
+        ("frealnvp", RealNVP),
+        ("spline", NeuralSplineFlow),
+        ("nsf", NeuralSplineFlow),
+        ("maf", MaskedAutoregressiveFlow),
+    ],
+)
+def test_get_native_flow_class(name, expected_class):
+    assert get_native_flow_class(name) is expected_class
+
+
+def test_get_native_flow_class_error():
+    with pytest.raises(ValueError, match=r"Unknown flow: invalid"):
+        get_native_flow_class("invalid")
+
+
+def test_get_flow_class_glasflow():
+    expected = object()
+    with patch(
+        "nessai.experimental.flows.glasflow.get_glasflow_class",
+        return_value=expected,
+    ) as mock_get:
+        out = get_flow_class("glasflow.realnvp")
+    mock_get.assert_called_once_with("glasflow.realnvp")
+    assert out is expected
+
+
+def test_get_flow_class_native():
+    expected = object()
+    with patch(
+        "nessai.flows.utils.get_native_flow_class",
+        return_value=expected,
+    ) as mock_get:
+        out = get_flow_class("realnvp")
+    mock_get.assert_called_once_with("realnvp")
+    assert out is expected
+
+
 def test_configure_model_basic(config):
     """Test configure model with the most basic config."""
     config["kwargs"] = dict(num_bins=2)
@@ -211,27 +259,12 @@ def test_configure_model_basic(config):
     )
 
 
-@pytest.mark.parametrize(
-    "flow_inputs",
-    [
-        {"ftype": "realnvp", "expected": "realnvp.RealNVP"},
-        {"ftype": "frealnvp", "expected": "realnvp.RealNVP"},
-        {"ftype": "spline", "expected": "nsf.NeuralSplineFlow"},
-        {"ftype": "nsf", "expected": "nsf.NeuralSplineFlow"},
-        {"ftype": "maf", "expected": "maf.MaskedAutoregressiveFlow"},
-    ],
-)
-def test_configure_model_flows(config, flow_inputs):
+def test_configure_model_ftype(config):
     """Test the different flows."""
-    config["ftype"] = flow_inputs["ftype"]
-    with patch(f"nessai.flows.{flow_inputs['expected']}") as mock_flow:
+    config["ftype"] = "realnvp"
+    with patch("nessai.flows.utils.get_native_flow_class") as mock_get:
         model, _ = configure_model(config)
-    mock_flow.assert_called_with(
-        config["n_inputs"],
-        config["n_neurons"],
-        config["n_blocks"],
-        config["n_layers"],
-    )
+    mock_get.assert_called_with("realnvp")
 
 
 def test_configure_model_flow_class(config):