update file structure

__init__.py → flax_models/__init__.py

@@ -1,18 +1,99 @@
-from .mlp_pinn import MLP_PINN
-from .PirateNet import PirateNet
 from .mlp import MLP
 from .siren import SIREN
 from .wire import WIRE
 from .activations import get_activation, list_activations
+from flax import serialization
+import os
+import yaml
+import jax
+import jax.numpy as jnp
+from ml_collections import ConfigDict
 
 model_key_dict = {
     "MLP": MLP,
     "SIREN": SIREN,
-    "WIRE": WIRE,
-    "PirateNet": PirateNet,
-    "MLP_PINN": MLP_PINN
+    "WIRE": WIRE
 }
 
+def make_model(config):
+    """
+    Create and configure a flax neural network nn.Module based on configuration.
+
+    Args:
+        config: Model configuration containing:
+            - model_name: Type of model (MLP, SIREN, WIRE, etc.)
+            - output_dim: Number of output dimensions
+            - hidden_dim: Hidden layer dimensions
+            - num_layers: Number of layers
+            - activation: Activation function name
+            - extra_model_args: Additional model-specific arguments
+            
+    Returns:
+        model (nn.Module): Configured flax nn.Module instance ready for training
+        
+    Note:
+        Handles special case for WIRE and SIREN models which don't accept
+        activation functions as an argument.
+    """
+
+    model = get_model(config.model_name)
+    if config.extra_model_args is not None:
+        if config.model_name == "WIRE" or config.model_name == "SIREN":
+            model = model(output_dim=config.output_dim,
+                          hidden_dim=config.hidden_dim,
+                          num_layers=config.num_layers,
+                          **config.extra_model_args)
+        else:
+            model = model(output_dim=config.output_dim,
+                        hidden_dim=config.hidden_dim, 
+                        num_layers=config.num_layers,
+                        act=get_activation(config.activation),
+                        **config.extra_model_args)
+    else:
+        model = model(output_dim=config.output_dim, 
+                      hidden_dim=config.hidden_dim, 
+                      num_layers=config.num_layers,
+                      act=get_activation(config.activation),
+    )   
+        
+    return model
+
+def load_metric_from_model(model_dir):
+    """
+    Load the model state from a given directory.
+    If the model has output dimension of 10, meaning
+    it was trained only on the symmetric part of the metric,
+    it reconstructs the full metric tensor.
+    
+    Args:
+        model_dir (str): Directory containing the model state file.
+
+    Returns:
+        callable: The metric tensor function from the model.
+
+    """
+    with open(os.path.join(model_dir, "params.msgpack"), "rb") as f:
+        params = serialization.msgpack_restore(f.read())
+    
+    with open(os.path.join(model_dir, "architecture.yml"), "r") as f:
+        config_model = yaml.load(f, Loader=yaml.FullLoader)
+
+    config_model = ConfigDict(config_model)
+    model = make_model(config_model.architecture)
+
+    if config_model.architecture.output_dim == 16:
+        return lambda coords: model.apply(params, coords).reshape(4, 4)
+    elif config_model.architecture.output_dim == 10:
+        return lambda coords: reconstruct_full_metric(model.apply(params, coords)).reshape(4, 4)
+
+def reconstruct_full_metric(metric_sym: jax.Array, n : int) -> jax.Array: 
+    """returns the fully reconstructed (n, n) metric tensor from the symmetry reduced metric""" 
+    i, j = jnp.triu_indices(n, k=0)
+    matrix = jnp.zeros((n, n))
+    matrix = matrix.at[i, j].set(metric_sym)
+    matrix = matrix.at[j, i].set(metric_sym)
+    return matrix
+
 def get_model(model_name : str):
     """
     Get the model class by name.
@@ -44,31 +125,7 @@ def create_model_configs():
             "first_omega_0": 4.,
             "hidden_omega_0": 4.,
             "scale": 5.,
-        },
-        "PirateNet": {
-            "nonlinearity": 0.0,
-            "pi_init": None,
-            "reparam": {
-                "type": "weight_fact",
-                "mean": 1.0,
-                "stddev": 0.1,
-            },
-            "fourier_emb": {
-                "embed_scale": 2.,
-                "embed_dim": 256,
-            },
-        },
-        "MLP_PINN": {
-            "reparam": {
-                "type": "weight_fact",
-                "mean": 1.0,
-                "stddev": 0.1,
-            },
-            "fourier_emb": {
-                "embed_scale": 2.,
-                "embed_dim": 256,
-            },
-        },
+        }
 
     }
     return model_configs

gw/cartesian/silu/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 128
+  model_name: MLP
+  num_layers: 5
+  output_dim: 16
+training:
+  metric_type: distortion

gw/cartesian/silu/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a71c9e504d93ddbb04ac9c9d325d1b93e18d9d46e33b212a9d17c3157ef2074
+size 341507

gw/cartesian/silu/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- true
+grid_range:
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+grid_shape:
+- 140
+- 10
+- 10
+- 140

gw/cartesian/siren/architecture.yml

@@ -0,0 +1,10 @@
+architecture:
+  activation: silu
+  extra_model_args:
+    omega_0: 1.0
+  hidden_dim: 128
+  model_name: SIREN
+  num_layers: 5
+  output_dim: 16
+training:
+  metric_type: distortion

gw/cartesian/siren/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea59e055998adf61be227c9c1b9a011348e6be52a621c353061c2567c0cd018d
+size 341591

gw/cartesian/siren/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- true
+grid_range:
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+grid_shape:
+- 140
+- 10
+- 10
+- 140

gw/cartesian/wire/architecture.yml

@@ -0,0 +1,12 @@
+architecture:
+  activation: silu
+  extra_model_args:
+    first_omega_0: 1
+    hidden_omega_0: 1
+    scale: 1
+  hidden_dim: 90
+  model_name: WIRE
+  num_layers: 5
+  output_dim: 16
+training:
+  metric_type: distortion

gw/cartesian/wire/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0d95701e9c4d6e037afbac8021d9fb9b0388b65180fb10ec0a2089fadc9b3b7
+size 337857

gw/cartesian/wire/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- true
+grid_range:
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+- - 0.0
+  - 10.0
+grid_shape:
+- 140
+- 10
+- 10
+- 140

kerr/boyer_lindquist/a_0.623/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 190
+  model_name: MLP
+  num_layers: 6
+  output_dim: 16
+training:
+  metric_type: distortion

kerr/boyer_lindquist/a_0.623/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed0b22feeba09f6106b75fdcc160c0abf76e7e654c8aa0986f6f163b442e4dcc
+size 742275

kerr/boyer_lindquist/a_0.623/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 3.0
+  - 8.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128

kerr/boyer_lindquist/a_0.628/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 190
+  model_name: MLP
+  num_layers: 6
+  output_dim: 16
+training:
+  metric_type: distortion

kerr/boyer_lindquist/a_0.628/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5be530410902ee569bbccadadef0827b022196b413812cc9d77564d01ebcb0f1
+size 742275

kerr/boyer_lindquist/a_0.628/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 3.0
+  - 8.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128

kerr/boyer_lindquist/prograde/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 190
+  model_name: MLP
+  num_layers: 6
+  output_dim: 16
+training:
+  metric_type: distortion

kerr/boyer_lindquist/prograde/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ac2705cc1280e5f97246ca752cf87384b1fc76ca2329259dfd540151ffec27a
+size 742275

kerr/boyer_lindquist/prograde/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 3.0
+  - 8.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128

kerr/boyer_lindquist/zackiger/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 190
+  model_name: MLP
+  num_layers: 6
+  output_dim: 16
+training:
+  metric_type: distortion

kerr/boyer_lindquist/zackiger/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82cecefe0018bb127fb5c7950ee7569793fd246736ee93839fef1bcf143a1e27
+size 742275

kerr/boyer_lindquist/zackiger/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 6.0
+  - 14.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128

kerr/kerr_schild_cartesian/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 190
+  model_name: MLP
+  num_layers: 5
+  output_dim: 16
+training:
+  metric_type: full_flatten

kerr/kerr_schild_cartesian/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9b7a968b6f2557abfad7e0f7267cbe20a2a57ef9d72fbf3c63a28b595663d1e
+size 742275

kerr/kerr_schild_cartesian/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- true
+grid_range:
+- - 0.0
+  - 0.0
+- - -3.0
+  - 3.0
+- - -3.0
+  - 3.0
+- - 0.1
+  - 3.0
+grid_shape:
+- 1
+- 128
+- 128
+- 128

schwarzschild/spherical/close_event_horizon/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 256
+  model_name: MLP
+  num_layers: 5
+  output_dim: 16
+training:
+  metric_type: distortion

schwarzschild/spherical/close_event_horizon/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a22f2792349d0281035e6eed84294d7422f7391f99d196ab0e3aa1a0b876cd8c
+size 1337877

schwarzschild/spherical/close_event_horizon/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 2.1
+  - 50.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128

schwarzschild/spherical/perihelion/architecture.yml

@@ -0,0 +1,9 @@
+architecture:
+  activation: silu
+  extra_model_args: {}
+  hidden_dim: 128
+  model_name: MLP
+  num_layers: 6
+  output_dim: 16
+training:
+  metric_type: distortion

schwarzschild/spherical/perihelion/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63a34941f366b9424b90f4807310aa937c966bcdbd781d027c3dfb4f8df309b1
+size 407620

schwarzschild/spherical/perihelion/train_data.yml

@@ -0,0 +1,19 @@
+endpoint:
+- true
+- true
+- true
+- false
+grid_range:
+- - 0.0
+  - 0.0
+- - 5.0
+  - 140.0
+- - 0.01
+  - 3.1315926535897933
+- - 0.0
+  - 6.283185307179586
+grid_shape:
+- 1
+- 128
+- 128
+- 128