Numerics

.github/workflows/build.yml

@@ -21,24 +21,28 @@ jobs:
 
     # Skip CI if 'skip ci' is contained in latest commit message
     if: "!contains(github.event.head_commit.message, 'skip ci')"
-    
+
     steps:
-    
+
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-    
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install .[dev]
-
+
+    - name: Run pre-test scripts
+      run: |
+        python tests/gen_test_data.py
+
     - name: Test with pytest
       run: |
         pytest
-    
+
     - name: Upload coverage to Codecov  
       uses: codecov/codecov-action@v3
       with:

careless/io/manager.py

@@ -453,6 +453,8 @@ def build_model(self, parser=None, surrogate_posterior=None, prior=None, likelih
             else:
                 raise ValueError(f"Unsupported scale bijector type, {parser.scale_bijector}")
 
+            istd = BaseModel.get_intensities(self.inputs).std()
+
             if parser.image_layers > 0:
                 from careless.models.scaling.image import NeuralImageScaler
                 n_images = np.max(BaseModel.get_image_id(self.inputs)) + 1
@@ -462,10 +464,14 @@ def build_model(self, parser=None, surrogate_posterior=None, prior=None, likelih
                     parser.mlp_layers,
                     mlp_width,
                     epsilon=parser.epsilon,
-                    scale_bijector=scale_bijector
+                    scale_bijector=scale_bijector,
+                    scale_multiplier=istd,
                 )
             else:
-                mlp_scaler = MLPScaler(parser.mlp_layers, mlp_width, epsilon=parser.epsilon, scale_bijector=scale_bijector)
+                mlp_scaler = MLPScaler(
+                    parser.mlp_layers, mlp_width, 
+                    epsilon=parser.epsilon, scale_bijector=scale_bijector, scale_multiplier=istd,
+                )
                 if parser.use_image_scales:
                     n_images = np.max(BaseModel.get_image_id(self.inputs)) + 1
                     image_scaler = ImageScaler(n_images)

careless/models/scaling/image.py

@@ -96,7 +96,7 @@ def call(self, metadata_and_image_id, *args, **kwargs):
         return result
 
 class NeuralImageScaler(Scaler):
-    def __init__(self, image_layers, max_images, mlp_layers, mlp_width, leakiness=0.01, epsilon=1e-7, scale_bijector=None):
+    def __init__(self, image_layers, max_images, mlp_layers, mlp_width, leakiness=0.01, epsilon=1e-7, scale_bijector=None, scale_multiplier=None):
         super().__init__()
         layers = []
         if leakiness is None:
@@ -111,7 +111,7 @@ def __init__(self, image_layers, max_images, mlp_layers, mlp_width, leakiness=0.
 
         self.image_layers = layers
         from careless.models.scaling.nn import MetadataScaler
-        self.metadata_scaler = MetadataScaler(mlp_layers, mlp_width, leakiness, epsilon=epsilon, scale_bijector=scale_bijector)
+        self.metadata_scaler = MetadataScaler(mlp_layers, mlp_width, leakiness, epsilon=epsilon, scale_bijector=scale_bijector, scale_multiplier=scale_multiplier)
 
     def call(self, inputs):
         result = self.get_metadata(inputs)

careless/models/scaling/nn.py

@@ -29,7 +29,7 @@ class MetadataScaler(Scaler):
     Neural network based scaler with simple dense layers.
     This neural network outputs a normal distribution.
     """
-    def __init__(self, n_layers, width, leakiness=0.01, epsilon=1e-7, scale_bijector=None):
+    def __init__(self, n_layers, width, leakiness=0.01, epsilon=1e-7, scale_bijector=None, scale_multiplier=None):
         """
         Parameters
         ----------
@@ -40,6 +40,13 @@ def __init__(self, n_layers, width, leakiness=0.01, epsilon=1e-7, scale_bijector
         leakiness : float or None
             If float, use LeakyReLU activation with provided parameter. Otherwise 
             use a simple ReLU
+        epsilon : float
+            A small constant for numerical stability. This is passed to the distribution layer.
+        scale_bijector : tfp.bijectors.Bijector
+            Optional scale bijector for the ouptut distibution
+        scale_multiplier : float
+            Optional constant to multiply the output location and scale by. This can increase
+            numerical stability. 
         """
         super().__init__()
 
@@ -74,6 +81,10 @@ def __init__(self, n_layers, width, leakiness=0.01, epsilon=1e-7, scale_bijector
         #The final layer converts the output to a Normal distribution
         #tfp_layers.append(tfp.layers.IndependentNormal())
         tfp_layers.append(NormalLayer(epsilon=epsilon, scale_bijector=scale_bijector))
+        if scale_multiplier is not None:
+            tfp_layers.append(
+                tfk.layers.Lambda(lambda x: tfb.Shift(scale_multiplier)(x))
+            )
 
         self.network = tfk.Sequential(mlp_layers)
         self.distribution = tfk.Sequential(tfp_layers)

pyproject.toml

@@ -55,3 +55,4 @@ testpaths = [
 
 [tool.setuptools.dynamic]
 version = {file = "careless/VERSION"}
+

tests/conftest.py

@@ -8,28 +8,23 @@
 import gemmi
 
 def pytest_sessionstart(session):
-    rundir = "data/"
-    rundir = abspath(join(dirname(__file__), rundir))
-
-    command = """
-    careless poly 
-        --disable-progress-bar 
-        --iterations=10 
-        --merge-half-datasets 
-        --half-dataset-repeats=3 
-        --test-fraction=0.1 
-        --disable-gpu 
-        --anomalous 
-        --wavelength-key=Wavelength
-        dHKL,Hobs,Kobs,Lobs,Wavelength
-        pyp_off.mtz 
-        pyp_2ms.mtz 
-        output/pyp
+    out_dir = "data/output"
+    out_dir = abspath(join(dirname(__file__), out_dir))
+    if exists(out_dir):
+        return
+    msg = f"""
+No test data in {out_dir}. Please run gen_test_data.py and retry.
+It is required to run this script prior to calling pytest. Here is
+an example of how to run the careless tests. 
+
+```
+cd careless
+pip install -e .[dev]
+python tests/gen_test_data.py
+pytest
+```
     """
-    if not exists(f"{rundir}/output"):
-        mkdir(f"{rundir}/output")
-        from subprocess import call
-        call(command.split(), cwd=rundir)
+    raise FileNotFoundError(msg)
 
 @pytest.fixture
 def cell_and_spacegroups():

tests/gen_test_data.py

@@ -0,0 +1,50 @@
+"""
+Generate the test data required for running the careless test suite. 
+It is required to run this script prior to calling pytest. Here is
+an example of how to run the careless tests. 
+
+```
+cd careless
+pip install -e .[dev]
+python tests/gen_test_data.py
+pytest
+```
+
+"""
+
+from os import listdir,mkdir
+from os.path import dirname, abspath, join, exists
+import numpy as np
+import pandas as pd
+import re
+import reciprocalspaceship as rs
+import gemmi
+
+
+
+def main():
+    rundir = "data/"
+    rundir = abspath(join(dirname(__file__), rundir))
+
+    command = """
+    careless poly 
+        --disable-progress-bar 
+        --iterations=10 
+        --merge-half-datasets 
+        --half-dataset-repeats=3 
+        --test-fraction=0.1 
+        --disable-gpu 
+        --anomalous 
+        --wavelength-key=Wavelength
+        dHKL,Hobs,Kobs,Lobs,Wavelength
+        pyp_off.mtz 
+        pyp_2ms.mtz 
+        output/pyp
+    """
+    if not exists(f"{rundir}/output"):
+        mkdir(f"{rundir}/output")
+    from subprocess import call
+    call(command.split(), cwd=rundir)
+
+if __name__=="__main__":
+    main()