subhaloscript 1.0.2__tar.gz → 1.0.3__tar.gz

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: subhaloscript
-Version: 1.0.2
+Version: 1.0.3
 Summary: Utility functions for analyzing subhalo distributions.
 Author-email: Charles Gannon <cgannon@ucmerced.edu>
 License-File: LICENSE

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/meta.yaml

@@ -1,6 +1,6 @@
 package:
   name: subhaloscript
-  version: "1.0.2"
+  version: "1.0.3"
 
 source:
   path: .  # Or use `git_url`/`url` if not building from local files

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 dependencies = ["numpy","pandas","scipy","h5py", "scikit-learn"]
 name = "subhaloscript"
-version = "1.0.2"
+version = "1.0.3"
 authors = [
   { name="Charles Gannon", email="cgannon@ucmerced.edu" },
 ]

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/subscript/macros.py

@@ -10,7 +10,7 @@ from numpy.dtypes import StringDType
 
 def macro_add(macros:dict[str, Callable], macro, label=None, **kwargs):
     """
-    Add a new macro function to an existing macros dictionary.
+    Add a new "script" to an existing macros dictionary.
 
     Parameters
     ----------

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/subscript/wrappers.py

@@ -6,6 +6,7 @@ from functools import reduce
 import numpy as np
 from numpy.typing import ArrayLike
 import h5py
+from copy import copy
 
 from subscript.util import is_arraylike
 from subscript.tabulatehdf5 import NodeProperties, tabulate_trees
@@ -27,6 +28,7 @@ def reduce_input(l, out=None):
     -------
     list of dict or UserDict
         Flattened list containing all dictionaries/UserDicts found in `l`.
+
     """
     if out is None:
         out = []
@@ -70,6 +72,32 @@ def format_nodedata(gout, out_index=-1)->Iterable[NodeProperties]:
         raise RuntimeError(f"Unrecognized data type for gout {type(gout)}")
     return _gout
 
+def _summarize(outs, summarize, statfuncs):
+    if summarize is None or not summarize:
+        return outs
+
+    _statfuncs = [np.mean, ] if statfuncs is None else statfuncs
+
+    if isinstance(outs[0], Iterable):
+        eval_stats = lambda f,m: f(np.asarray([treeo[m] for treeo in outs]), axis=0)
+        summary = [[eval_stats(f,m) for m, _ in enumerate(outs[0])] for f in _statfuncs]
+    else:
+        eval_stats = lambda f: f(np.asarray([treeo for treeo in outs]), axis=0)
+        summary = [eval_stats(f) for f in _statfuncs]
+    return summary
+
+
+# Eliminate lists of 1 item recursively
+def _format_out(o):
+    if (not isinstance(o, Iterable)) or (isinstance(o, str)):
+        return o
+    if len(o) == 1:
+        return _format_out(o[0])
+    out = [_format_out(i) for i in o]
+    if isinstance(o, np.ndarray):
+        return np.asarray(out)
+    return out
+
 def gscript(func):
     """
     Decorator to wrap a function that processes Galacticus-like node data.
@@ -150,30 +178,9 @@ def gscript(func):
             _o = [o,] if single_out else o
             outs.append(_o)
 
-        # Eliminate lists of 1 item recursively
-        def format_out(o):
-            if (not isinstance(o, Iterable)) or (isinstance(o, str)):
-                return o
-            if len(o) == 1:
-                return format_out(o[0])
-            out = [format_out(i) for i in o]         
-            if isinstance(o, np.ndarray):
-                return np.asarray(out)
-            return out
-
-        if not summarize:
-            return format_out(outs)
-    
-        _statfuncs = [np.mean, ] if statfuncs is None else statfuncs
-
-        if isinstance(outs[0], Iterable):
-            eval_stats = lambda f,m: f(np.asarray([treeo[m] for treeo in outs]), axis=0)
-            summary = [[eval_stats(f,m) for m, _ in enumerate(outs[0])] for f in _statfuncs] 
-        else:
-            eval_stats = lambda f: f(np.asarray([treeo for treeo in outs]), axis=0)
-            summary = [eval_stats(f) for f in _statfuncs] 
+        summary = _summarize(outs, summarize=summarize, statfuncs=statfuncs)
 
-        return format_out(summary)
+        return _format_out(summary)
     return wrap
 
 def gscript_proj(func):
@@ -208,27 +215,36 @@ def gscript_proj(func):
     """
     def wrap(gout, normvector, *args, **kwargs):
         normvector = np.asarray(normvector)
-        n = None
+
+        if normvector.ndim == 1:
+            #raise NotImplementedError()
+            return gscript(func)(gout, *args, normvector=normvector, **kwargs)
+        if normvector.ndim > 2 or normvector.ndim <= 0:
+            raise RuntimeError(f'"normvector" must be either 1 or 2 dimensional')
 
         @gscript
-        def wrap_inner(gout, *args, normvector, **kwargs):
-            nonlocal n
-            v =  normvector
-            if n is not None:
-                v = normvector[n]
-                n += 1  
-                if n >= len(normvector):
-                    n = 0
+        def wrap_inner_main(gout, **kwargs):
+            n = gout.unfilter()['__custom_proj_iter__'][0]
+            return func(gout, *args, normvector=normvector[n], **kwargs)
 
 
-            return func(gout, *args, normvector=v, **kwargs)
+        ## This wrapped as well so we can call with None
+        def wrap_inner(gout, **kwargs2):
+            _gout = format_nodedata(gout)[0].unfilter()
 
+            _input = []
+            for n, _ in enumerate(normvector):
+                _in = copy(_gout)
+                _in.data = copy(_in.data)
+                _in.data['__custom_proj_iter__'] = n * np.ones(_in.data[next(_in.data.__iter__())].shape[0], dtype=int)
+                _input.append(_in)
 
-        if normvector.ndim == 1:
-            return wrap_inner(gout, *args, normvector=normvector, **kwargs)
+            return wrap_inner_main(_input, **(kwargs | kwargs2))
 
-        n = 0
-        return wrap_inner([gout for _ in normvector], *args, normvector=normvector, **kwargs)
+        if gout is None:
+            return wrap_inner
+
+        return wrap_inner(gout)
 
     return wrap
 
@@ -267,5 +283,3 @@ def multiproj(func, nfilter):
         Function wrapped with `gscript_proj` and the node filter applied.
     """
     return gscript_proj(freeze(func, nfilter=nfilter))
-
-

{subhaloscript-1.0.2 → subhaloscript-1.0.3}/tests/test_wrappers.py

@@ -53,8 +53,12 @@ def test_multi_proj():
     out_expected_yz = np.asarray((4, 3, 0), dtype=int)
 
     out_expected = np.mean(np.asarray((out_expected_xy, out_expected_xz, out_expected_yz)), axis=0)
-        
-    out_actual, dn_r = spatial2d_dn(mockdata, bins=bins, normvector=norm, summarize=True)
+    print(out_expected)
+
+    out = spatial2d_dn(mockdata, bins=bins, normvector=norm, summarize=True)
+    print(out)
+    out_actual, dn_r = out
+
 
     testing.assert_allclose(dn_r, bins)
     testing.assert_allclose(out_actual, out_expected) 
@@ -96,6 +100,44 @@ def test_gscript_proj_wrap():
 
     testing.assert_allclose(n_actual, n_expected)
 
+
+def test_gscript_proj_wrap_none():
+    test_x       = np.asarray((0.0, 0.25, 0.5       , 0.7       , 0.8        , 1.3, 1.4))
+    test_y       = np.asarray((0.0, 0.00, 0.5       , 0.3       , 0.9        , 0.0, 0.0))
+    test_z       = np.asarray((0.0, 0.00, 0.5       , 0.4       , 0.1        , 0.0, 0.0))
+
+    mockdata = {
+                    ParamKeys.x: test_x,
+                    ParamKeys.y: test_y,
+                    ParamKeys.z: test_z
+    }
+
+    r_xy_expected = np.asarray((0.0, 0.25, 0.70710678, 0.76157731, 1.20415946, 1.3, 1.4))
+    r_xz_expected = np.asarray((0.0, 0.25, 0.70710678, 0.80622577, 0.80622577, 1.3, 1.4))
+    r_yz_expected = np.asarray((0.0, 0.00, 0.70710678, 0.5       , 0.90553851, 0.0, 0.0))
+
+    rmin, rmax = 0.2, 0.71
+    normvectors = np.identity(3)
+
+    n_expected = (
+                   np.sum((r_xy_expected >= rmin) & (r_xy_expected <= rmax))
+                  +np.sum((r_xz_expected >= rmin) & (r_xz_expected <= rmax))
+                  +np.sum((r_yz_expected >= rmin) & (r_yz_expected <= rmax))
+                ) / 3
+
+    nfproj   = freeze(r2d, rmin=rmin, rmax=rmax)
+
+    frozentest = gscript_proj(freeze(nodecount, nfilter=nfproj))(None, summarize=True, normvector=normvectors)
+    n_actual = frozentest(mockdata)
+
+    testing.assert_allclose(n_actual, n_expected)
+
+    normvectors2 = ((1,0,0), (0,1,0), (0,0,1))
+
+    n_actual = gscript_proj(freeze(nodecount, nfilter=nfproj))(mockdata, summarize=True, normvector=normvectors2)
+
+    testing.assert_allclose(n_actual, n_expected)
+
 def test_multiproj():
     test_x       = np.asarray((0.0, 0.25, 0.5       , 0.7       , 0.8        , 1.3, 1.4))
     test_y       = np.asarray((0.0, 0.00, 0.5       , 0.3       , 0.9        , 0.0, 0.0))
@@ -208,4 +250,5 @@ def test_autofreeze():
                 testing.assert_allclose(f[key][_key][:], _val)       
 
 
-
+if __name__ == "__main__":
+    test_multi_proj()