Arr API tweaks

CHANGELOG.rst

@@ -6,6 +6,12 @@ What's new
 
 All notable changes to the codebase are documented in this file. Changes that may result in differences in model output, or are required in order to run an old parameter set with the current version, are flagged with the term "Regression information".
 
+Version 0.5.3 (2024-06-10)
+--------------------------
+- ``ss.uids`` class implements set operators to facilitate combining or otherwise operating on collections of UIDs
+- ``FloatArr.isnan`` and ``FloatArr.notnan`` return ``BoolArr`` instances rather than UIDs (so as to facilitate logical operations with other ``BoolArr`` instances, and to align more closely with `np.isnan`)
+- ``Arr.true()`` and ``Arr.false()`` are supported for all ``Arr`` subclasses
+- ``BoolArr.isnan`` and ``Boolarr.notnan`` are also implemented (although since ``BoolArr`` cannot store NaN values, these always return ``False`` and ``True``, respectively)
 
 Version 0.5.2 (2024-06-04)
 --------------------------

starsim/network.py

@@ -347,11 +347,10 @@ def available(self, people, sex):
         # This property could also be overwritten by a NetworkConnector
         # which could incorporate information about membership in other
         # contact networks
-        right_sex = people[sex]
-        is_active = self.active(people)
-        is_available = (right_sex & is_active).uids
-        still_available = is_available.remove(self.members)
-        return still_available
+        available = people[sex] & self.active(people)
+        available[self.edges.p1] = False
+        available[self.edges.p2] = False
+        return available.uids
 
     def beta_per_dt(self, disease_beta=None, dt=None, uids=None):
         if uids is None: uids = Ellipsis

starsim/people.py

@@ -284,7 +284,7 @@ def remove_dead(self):
                 network.remove_uids(uids) # TODO: only run once every nth timestep
 
             # Calculate the indices to keep
-            self.auids = self.auids.remove(uids)
+            self.auids = self.auids[np.isin(self.auids, np.unique(uids), assume_unique=True, invert=True, kind='sort')]
 
         return
 

starsim/states.py

@@ -72,7 +72,7 @@ class Arr(np.lib.mixins.NDArrayOperatorsMixin):
         coerce (bool): Whether to ensure the the data is one of the supported data types
         skip_init (bool): Whether to skip initialization with the People object (used for uid and slot states)
     """
-    def __init__(self, name, dtype=None, default=None, nan=None, raw=None, label=None, coerce=True, skip_init=False):
+    def __init__(self, name, dtype=None, default=None, nan=None, raw=None, label=None, coerce=True, skip_init=False, people=None):
         if coerce:
             dtype = check_dtype(dtype, default)
 
@@ -85,7 +85,7 @@ def __init__(self, name, dtype=None, default=None, nan=None, raw=None, label=Non
 
         # Properties that are initialized later
         self.raw = np.empty(0, dtype=dtype)
-        self.people = None # Used solely for accessing people.auids
+        self.people = people # Used solely for accessing people.auids
         self.len_used = 0
         self.len_tot = 0
         self.initialized = skip_init
@@ -187,7 +187,15 @@ def set_nan(self, uids):
         """ Shortcut function to set values to NaN """
         self.raw[uids] = self.nan
         return
-
+
+    @property
+    def isnan(self):
+        return self.asnew(self.values == self.nan, cls=BoolArr)
+
+    @property
+    def notnan(self):
+        return self.asnew(self.values != self.nan, cls=BoolArr)
+
     def grow(self, new_uids=None, new_vals=None):
         """
         Add new agents to an Arr
@@ -246,22 +254,30 @@ def asnew(self, arr=None, cls=None):
         new.raw[new.auids] = arr
         return new
 
+    def true(self):
+        """ Efficiently convert truthy values to UIDs """
+        return self.auids[self.values.astype(bool)]
+
+    def false(self):
+        """ Reverse of true(); return UIDs of falsy values """
+        return self.auids[~self.values.astype(bool)]
+
 
 class FloatArr(Arr):
     """ Subclass of Arr with defaults for floats """
-    def __init__(self, name, default=None, nan=np.nan, label=None, skip_init=False):
-        super().__init__(name=name, dtype=ss_float, default=default, nan=nan, label=label, coerce=False, skip_init=skip_init)
+    def __init__(self, name, nan=np.nan, **kwargs):
+        super().__init__(name=name, dtype=ss_float, nan=nan, coerce=False, **kwargs)
         return
-    
+
     @property
     def isnan(self):
-        """ Return indices that are NaN """
-        return np.nonzero(np.isnan(self.values))[0]
+        """ Return BoolArr for NaN values """
+        return self.asnew(np.isnan(self.values), cls=BoolArr)
 
     @property
     def notnan(self):
-        """ Return indices that are not-NaN """
-        return np.nonzero(~np.isnan(self.values))[0]
+        """ Return BoolArr for non-NaN values """
+        return self.asnew(~np.isnan(self.values), cls=BoolArr)
 
     @property
     def notnanvals(self):
@@ -270,30 +286,31 @@ def notnanvals(self):
         out = vals[np.nonzero(~np.isnan(vals))[0]]
         return out
 
-    
+
 class BoolArr(Arr):
     """ Subclass of Arr with defaults for booleans """
-    def __init__(self, name, default=None, nan=False, label=None, skip_init=False): # No good NaN equivalent for bool arrays
-        super().__init__(name=name, dtype=ss_bool, default=default, nan=nan, label=label, coerce=False, skip_init=skip_init)
+    def __init__(self, name, nan=False, **kwargs): # No good NaN equivalent for bool arrays
+        super().__init__(name=name, dtype=ss_bool, nan=nan, coerce=False, **kwargs)
         return
 
     def __and__(self, other): return self.asnew(self.values & other)
     def __or__(self, other):  return self.asnew(self.values | other)
     def __xor__(self, other): return self.asnew(self.values ^ other)
     def __invert__(self):     return self.asnew(~self.values)
-
+
+    # BoolArr cannot store NaNs so report all entries as being not-NaN
     @property
-    def uids(self):
-        """ Efficiently convert True values to UIDs """
-        return self.auids[np.nonzero(self.values)[0]]
+    def isnan(self):
+        return self.asnew(np.full_like(self.values, fill_value=False), cls=BoolArr)
 
-    def true(self):
-        """ Alias to BoolArr.uids """
-        return self.uids
+    @property
+    def notnan(self):
+        return self.asnew(np.full_like(self.values, fill_value=True), cls=BoolArr)
 
-    def false(self):
-        """ Reverse of true(); return UIDs of values that are false """
-        return self.auids[np.nonzero(~self.values)[0]]
+    @property
+    def uids(self):
+        """ Alias to Arr.true """
+        return self.true()
 
     def split(self):
         """ Return UIDs of values that are true and false as separate arrays """
@@ -313,14 +330,7 @@ def __init__(self, name, label=None):
     def uids(self):
         """ Alias to self.values, to allow Arr.uids like BoolArr """
         return self.values
-
-    @property
-    def isnan(self):
-        return np.nonzero(self.values == self.nan)[0]
 
-    @property
-    def notnan(self):
-        return np.nonzero(self.values != self.nan)[0]
 
     def grow(self, new_uids=None, new_vals=None):
         """ Change the size of the array """
@@ -361,16 +371,27 @@ def cat(cls, *args, **kw):
 
     def remove(self, other, **kw):
         """ Remove provided UIDs from current array"""
-        return np.setdiff1d(self, other, assume_unique=True, **kw).view(self.__class__)
+        return np.setdiff1d(self, other, **kw).view(self.__class__)
 
     def intersect(self, other, **kw):
-        """ Keep only UIDs that match other array """
-        return np.intersect1d(self, other, assume_unique=True, **kw).view(self.__class__)
-
+        """ Keep only UIDs that are also present in the other array """
+        return np.intersect1d(self, other, **kw).view(self.__class__)
+
+    def union(self, other, **kw):
+        """ Return all UIDs present in both arrays """
+        return np.union1d(self, other, **kw).view(self.__class__)
+
     def to_numpy(self):
         """ Convert to a standard NumPy array """
         return np.array(self)
-
+
+    # Implement collection of operators
+    def __and__(self, other): return self.intersect(other)
+    def __or__(self, other):  return self.union(other)
+    def __sub__(self, other): return self.remove(other)
+    def __xor__(self, other): return np.setxor1d(self, other).view(self.__class__)
+    def __invert__(self):     raise Exception(f"Cannot invert an instance of {self.__class__.__name__}. One possible cause is attempting `~x.uids` - use `x.false()` or `(~x).uids` instead")
+
 
 class BooleanOperationError(NotImplementedError):
     """ Raised when a logical operation is performed on a non-logical array """

starsim/version.py

@@ -4,6 +4,6 @@
 
 __all__ = ['__version__', '__versiondate__', '__license__']
 
-__version__ = '0.5.2'
-__versiondate__ = '2024-06-04'
+__version__ = '0.5.3'
+__versiondate__ = '2024-06-10'
 __license__ = f'Starsim {__version__} ({__versiondate__}) — © 2023-2024 by IDM'