Repository: starsimhub/starsim (language: python)

Issue:
Check that dead-but-not-removed agents are being handled sensibly
              if we implement this change, wouldn't we need to go through all the modules and make sure that we were always conditioning on people being alive before selecting them for e.g. pair formation, disease acquisition/progression, etc?

_Originally posted by @robynstuart in https://github.com/amath-idm/starsim/issues/294#issuecomment-1958601482_
            
Rethink uidarray2
Addresses @RomeshA 's comments on #456, made as a separate PR to clarify changes.

Major change is that `IntArr` has been removed. In practice, this was only ever used for time index arrays (e.g. `ti_dead`). All the comparisons (e.g. `ti_dead <= sim.ti`) work just as well (even slightly faster) as floats, plus this lets us natively use `np.nan`. I also personally think we should consider moving away from a time index towards absolute time (see #482), but that's definitely beyond the scope of this PR!


Discussion / hints:
Agree with @cliffckerr that we should keep the option for users to turn off removing dead agents entirely, so we probably can't avoid having to go through all of the modules to check this. We might actually want to parametrize some of the pytest tests to test with/without removing dead agents
There are dead agents up to step 10:
```py
import starsim as ss

def check_dead(sim):
    print(sum(~sim.people.alive))
    
sim = ss.demo(analyzers=check_dead)
```
This **is** a bug: add these lines to the end of `Infection.make_new_cases()`:
```py
                assert sum(~sim.people.alive[trg]) == 0
                assert sum(~sim.people.alive[src]) == 0
```
Then rerun. Dead people can be sources of infection.


Fix patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index fd05fa86..87737bd5 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,18 @@ 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.4.0 (2024-04-24)
+--------------------------
+- Replace ``UIDArray``, ``ArrayView``, and ``State`` with ``Arr``, which has different subclasses for different data types (e.g. ``FloatArr``, ``BoolArr``, and ``IndexArr``). States are usually represented by ``BoolArr`` (e.g. ``sir.infected``), while other agent properties are represented by ``FloatArr`` (e.g. ``sir.rel_trans``).
+- Arrays that had previously been represented using an integer data type (e.g. ``sir.ti_infected``) are now also ``FloatArr``, to allow the use of ``np.nan``. Integer arrays are supported via ``IndexArr``, but these are only intended for use for slots and UIDs.
+- ``Arr`` objects automatically skip over dead (or otherwise removed) agents; the "active" UIDs are stored in ``sim.people.auids``, which is updated when agents are born or die. This array is linked to each ``Arr``, so that e.g. ``sim.people.age.mean()`` will only calculate the mean over alive agents. To access the underlying Numpy array, use ``sim.people.age.raw``.
+- ``FloatArr`` has ``isnan``, ``notnan``, and ``notnanvals`` properties. ``BoolArr`` has logical operations defined. For example, ``~people.female`` works, but ``~people.ti_dead`` does not; ``people.ti_dead.notnan`` works, but ``people.female.notnan`` does not.
+- UIDs used to be NumPy integer arrays; they are now ``ss.uids`` objects (which is a class, but is lowercase for consistency with ``np.array()``, which it is functionally similar to). Indexing a state by an integer array rather than ``ss.uids()`` now raises an exception, due to the ambiguity involved. To index the underlying array with an integer array, use ``Arr.raw[int_arr]``; to index only the active/alive agents, use ``Arr.values[int_arr]``.
+- Dead agents are no longer removed, so ``uid`` always corresponds to the position in the array. This means that no remapping is necessary, which has a significant performance benefit (roughly 2x faster for large numbers of agents).
+- Renamed ``omerge`` to ``dictmerge`` and ``omergeleft`` to ``dictmergeleft``.
+- *GitHub info*: PR `456 <https://github.com/starsimhub/starsim/pull/456>`_
+
+
 Version 0.3.4 (2024-04-18)
 --------------------------
 - Default duration of edges in ``ss.RandomNet`` changed from 1 to 0; this does not matter if ``dt=1``, but does matter with smaller ``dt`` values.
diff --git a/starsim/connectors.py b/starsim/connectors.py
index 4cdac665..089d8fe6 100644
--- a/starsim/connectors.py
+++ b/starsim/connectors.py
@@ -18,7 +18,7 @@ class simple_hiv_ng(Connector):
 
     def __init__(self, pars=None):
         super().__init__(pars=pars, label='HIV-Gonorrhea', diseases=[ss.HIV, ss.Gonorrhea])
-        self.pars = ss.omerge({
+        self.pars = ss.dictmerge({
             'rel_trans_hiv': 2,
             'rel_trans_aids': 5,
             'rel_sus_hiv': 2,
diff --git a/starsim/demographics.py b/starsim/demographics.py
index b1d42f35..cd132483 100644
--- a/starsim/demographics.py
+++ b/starsim/demographics.py
@@ -34,14 +34,14 @@ def __init__(self, pars=None, metadata=None, **kwargs):
         super().__init__(pars, **kwargs)
 
         # Set defaults
-        self.pars = ss.omergeleft(self.pars,
+        self.pars = ss.dictmergeleft(self.pars,
             birth_rate = 0,
             rel_birth = 1,
             units = 1e-3,  # assumes birth rates are per 1000. If using percentages, switch this to 1
         )
 
         # Process metadata. Defaults here are the labels used by UN data
-        self.metadata = ss.omergeleft(metadata,
+        self.metadata = ss.dictmergeleft(metadata,
             data_cols = dict(year='Year', cbr='CBR'),
         )
 
@@ -144,18 +144,18 @@ def __init__(self, pars=None, par_dists=None, metadata=None, **kwargs):
         """
         super().__init__(pars, **kwargs)
 
-        self.pars = ss.omergeleft(self.pars,
+        self.pars = ss.dictmergeleft(self.pars,
             rel_death = 1,
             death_rate = 20,  # Default = a fixed rate of 2%/year, overwritten if data provided
             units = 1e-3,  # assumes death rates are per 1000. If using percentages, switch this to 1
         )
 
-        self.par_dists = ss.omergeleft(par_dists,
+        self.par_dists = ss.dictmergeleft(par_dists,
             death_rate = ss.bernoulli
         )
 
         # Process metadata. Defaults here are the labels used by UN data
-        self.metadata = ss.omergeleft(metadata,
+        self.metadata = ss.dictmergeleft(metadata,
             data_cols = dict(year='Time', sex='Sex', age='AgeGrpStart', value='mx'),
             sex_keys = dict(f='Female', m='Male'),
         )
@@ -175,6 +175,7 @@ def make_death_prob_fn(module, sim, uids):
             death_rate = module.death_rate_data
 
         else:
+            ppl = sim.people
             data_cols = sc.objdict(module.metadata.data_cols)
             year_label = data_cols.year
             age_label  = data_cols.age
@@ -188,16 +189,20 @@ def make_death_prob_fn(module, sim, uids):
 
             df = module.death_rate_data.loc[module.death_rate_data[year_label] == nearest_year]
             age_bins = df[age_label].unique()
-            age_inds = np.digitize(sim.people.age[uids], age_bins) - 1
 
             f_arr = df[val_label].loc[df[sex_label] == sex_keys['f']].values
             m_arr = df[val_label].loc[df[sex_label] == sex_keys['m']].values
 
             # Initialize
             death_rate_df = pd.Series(index=uids)
-            death_rate_df[uids[sim.people.female[uids]]] = f_arr[age_inds[sim.people.female[uids]]] # TODO: avoid double indexing
-            death_rate_df[uids[sim.people.male[uids]]] = m_arr[age_inds[sim.people.male[uids]]]
-            death_rate_df[uids[sim.people.age[uids] < 0]] = 0  # Don't use background death rates for unborn babies
+            f_uids = uids.intersect(ppl.female.uids) # TODO: reduce duplication
+            m_uids = uids.intersect(ppl.male.uids)
+            f_age_inds = np.digitize(ppl.age[f_uids], age_bins) - 1
+            m_age_inds = np.digitize(ppl.age[m_uids], age_bins) - 1
+            death_rate_df[f_uids] = f_arr[f_age_inds]
+            death_rate_df[m_uids] = m_arr[m_age_inds]
+            unborn_inds = uids.intersect((sim.people.age < 0).uids)
+            death_rate_df[unborn_inds] = 0  # Don't use background death rates for unborn babies
 
             death_rate = death_rate_df.values
 
@@ -227,8 +232,7 @@ def update(self, sim):
 
     def apply_deaths(self, sim):
         """ Select people to die """
-        alive_uids = ss.true(sim.people.alive)
-        death_uids = self.pars.death_rate.filter(alive_uids)
+        death_uids = self.pars.death_rate.filter()
         sim.people.request_death(death_uids)
         return len(death_uids)
 
@@ -250,17 +254,17 @@ def __init__(self, pars=None, par_dists=None, metadata=None, **kwargs):
 
         # Other, e.g. postpartum, on contraception...
         self.add_states(
-            ss.State('infertile', bool, False),  # Applies to girls and women outside the fertility window
-            ss.State('fecund', bool, True),  # Applies to girls and women inside the fertility window
-            ss.State('pregnant', bool, False),  # Currently pregnant
-            ss.State('postpartum', bool, False),  # Currently post-partum
-            ss.State('ti_pregnant', int, ss.INT_NAN),  # Time pregnancy begins
-            ss.State('ti_delivery', int, ss.INT_NAN),  # Time of delivery
-            ss.State('ti_postpartum', int, ss.INT_NAN),  # Time postpartum ends
-            ss.State('ti_dead', int, ss.INT_NAN),  # Maternal mortality
+            ss.BoolArr('infertile'),  # Applies to girls and women outside the fertility window
+            ss.BoolArr('fecund', default=True),  # Applies to girls and women inside the fertility window
+            ss.BoolArr('pregnant'),  # Currently pregnant
+            ss.BoolArr('postpartum'),  # Currently post-partum
+            ss.FloatArr('ti_pregnant'),  # Time pregnancy begins
+            ss.FloatArr('ti_delivery'),  # Time of delivery
+            ss.FloatArr('ti_postpartum'),  # Time postpartum ends
+            ss.FloatArr('ti_dead'),  # Maternal mortality
         )
 
-        self.pars = ss.omergeleft(self.pars,
+        self.pars = ss.dictmergeleft(self.pars,
             dur_pregnancy = 0.75,
             dur_postpartum = 0.5,
             fertility_rate = 0,    # Usually this will be provided in CSV format
@@ -270,14 +274,14 @@ def __init__(self, pars=None, par_dists=None, metadata=None, **kwargs):
             units = 1e-3,          # Assumes fertility rates are per 1000. If using percentages, switch this to 1
         )
 
-        self.par_dists = ss.omergeleft(par_dists,
+        self.par_dists = ss.dictmergeleft(par_dists,
             fertility_rate = ss.bernoulli,
             maternal_death_rate = ss.bernoulli,
             sex_ratio = ss.bernoulli
         )
 
         # Process metadata. Defaults here are the labels used by UN data
-        self.metadata = ss.omergeleft(metadata,
+        self.metadata = ss.dictmergeleft(metadata,
             data_cols = dict(year='Time', age='AgeGrp', value='ASFR'),
         )
 
@@ -320,7 +324,7 @@ def make_fertility_prob_fn(module, sim, uids):
 
             # Adjust rates: rates are based on the entire population, but we need to remove
             # anyone already pregnant and then inflate the rates for the remainder
-            pregnant_uids = ss.true(module.pregnant[uids])  # Find agents who are already pregnant
+            pregnant_uids = module.pregnant.uids # Find agents who are already pregnant
             pregnant_age_counts, _ = np.histogram(sim.people.age[pregnant_uids], age_bins)  # Count them by age
             age_counts, _ = np.histogram(sim.people.age[uids], age_bins)  # Count overall number per age bin
             new_denom = age_counts - pregnant_age_counts  # New denominator for rates
@@ -392,7 +396,7 @@ def update_states(self, sim):
         self.fecund[postpartum] = True
 
         # Maternal deaths
-        maternal_deaths = ss.true(self.ti_dead <= sim.ti)
+        maternal_deaths = (self.ti_dead <= sim.ti).uids
         sim.people.request_death(maternal_deaths)
 
         return
@@ -401,14 +405,10 @@ def make_pregnancies(self, sim):
         """
         Select people to make pregnant using incidence data
         """
-        # Abbreviate
-        ppl = sim.people
-
         # People eligible to become pregnant. We don't remove pregnant people here, these
         # are instead handled in the fertility_dist logic as the rates need to be adjusted
-        denom_conds = ppl.female & ppl.alive
-        inds_to_choose_from = ss.true(denom_conds)
-        conceive_uids = self.pars.fertility_rate.filter(inds_to_choose_from)
+        eligible_uids = sim.people.female.uids
+        conceive_uids = self.pars.fertility_rate.filter(eligible_uids)
 
         # Set prognoses for the pregnancies
         if len(conceive_uids) > 0:
diff --git a/starsim/disease.py b/starsim/disease.py
index 17edb8ad..2af7c713 100644
--- a/starsim/disease.py
+++ b/starsim/disease.py
@@ -74,7 +74,7 @@ def set_initial_states(self, sim):
         Set initial values for states
 
         This could involve passing in a full set of initial conditions,
-        or using init_prev, or other. Note that this is different to initialization of the State objects
+        or using init_prev, or other. Note that this is different to initialization of the Arr objects
         i.e., creating their dynamic array, linking them to a People instance. That should have already
         taken place by the time this method is called. This method is about supplying initial values
         for the states (e.g., seeding initial infections)
@@ -180,11 +180,11 @@ class Infection(Disease):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.add_states(
-            ss.State('susceptible', bool, True),
-            ss.State('infected', bool, False),
-            ss.State('rel_sus', float, 1.0),
-            ss.State('rel_trans', float, 1.0),
-            ss.State('ti_infected', int, ss.INT_NAN),
+            ss.BoolArr('susceptible', default=True),
+            ss.BoolArr('infected'),
+            ss.FloatArr('rel_sus', default=1.0),
+            ss.FloatArr('rel_trans', default=1.0),
+            ss.FloatArr('ti_infected'),
         )
 
         self.rng_target = ss.random(name='target')
@@ -203,15 +203,14 @@ def infectious(self):
     def set_initial_states(self, sim):
         """
         Set initial values for states. This could involve passing in a full set of initial conditions,
-        or using init_prev, or other. Note that this is different to initialization of the State objects
+        or using init_prev, or other. Note that this is different to initialization of the Arr objects
         i.e., creating their dynamic array, linking them to a People instance. That should have already
         taken place by the time this method is called.
         """
         if self.pars.init_prev is None:
             return
 
-        alive_uids = ss.true(sim.people.alive)  # Maybe just sim.people.uid?
-        initial_cases = self.pars.init_prev.filter(alive_uids)
+        initial_cases = self.pars.init_prev.filter()
         self.set_prognoses(sim, initial_cases)  # TODO: sentinel value to indicate seeds?
         return
 
@@ -270,8 +269,8 @@ def make_new_cases(self, sim):
 
             nbetas = betamap[nkey]
             contacts = net.contacts
-            rel_trans = (self.infectious & people.alive) * self.rel_trans
-            rel_sus = (self.susceptible & people.alive) * self.rel_sus
+            rel_trans = self.rel_trans.asnew(self.infectious * self.rel_trans)
+            rel_sus   = self.rel_sus.asnew(self.susceptible * self.rel_sus)
             p1p2b0 = [contacts.p1, contacts.p2, nbetas[0]]
             p2p1b1 = [contacts.p2, contacts.p1, nbetas[1]]
             for src, trg, beta in [p1p2b0, p2p1b1]:
@@ -297,8 +296,8 @@ def make_new_cases(self, sim):
                 
         # Tidy up
         if len(new_cases) and len(sources):
-            new_cases = np.concatenate(new_cases)
-            sources = np.concatenate(sources)
+            new_cases = ss.uids.cat(new_cases)
+            sources = ss.uids.cat(sources)
         else:
             new_cases = np.empty(0, dtype=int)
             sources = np.empty(0, dtype=int)
@@ -310,7 +309,7 @@ def make_new_cases(self, sim):
 
     def _set_cases(self, sim, target_uids, source_uids=None):
         congenital = sim.people.age[target_uids] <= 0
-        if len(ss.true(congenital)) > 0:
+        if np.count_nonzero(congenital):
             src_c = source_uids[congenital] if source_uids is not None else None
             self.set_congenital(sim, target_uids[congenital], src_c)
         src_p = source_uids[~congenital] if source_uids is not None else None
diff --git a/starsim/diseases/cholera.py b/starsim/diseases/cholera.py
index 5bed389f..a5bed241 100644
--- a/starsim/diseases/cholera.py
+++ b/starsim/diseases/cholera.py
@@ -17,7 +17,7 @@ class Cholera(ss.Infection):
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         """ Initialize with parameters """
 
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             # Natural history parameters, all specified in days
             dur_exp2inf = ss.lognorm_ex(mean=2.772, stdev=4.737),  # Calculated from Azman et al. estimates https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3677557/
             dur_asymp2rec = ss.uniform(low=1, high=10),    # From WHO cholera fact sheet, asymptomatic individuals shed bacteria for 1-10 days (https://www.who.int/news-room/fact-sheets/detail/cholera)
@@ -39,7 +39,7 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
             p_env_transmit = 0,    # Probability of environmental transmission - filled out later
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_exp2inf    = ss.lognorm_ex,
             dur_asymp2rec  = ss.uniform,
             dur_symp2rec   = ss.lognorm_ex,
@@ -56,15 +56,15 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         
         self.add_states(
             # Susceptible & infected are added automatically, here we add the rest
-            ss.State('exposed', bool, False),
-            ss.State('symptomatic', bool, False),
-            ss.State('recovered', bool, False),
+            ss.BoolArr('exposed'),
+            ss.BoolArr('symptomatic'),
+            ss.BoolArr('recovered'),
     
             # Timepoint states
-            ss.State('ti_exposed', float, np.nan),
-            ss.State('ti_symptomatic', float, np.nan),
-            ss.State('ti_recovered', float, np.nan),
-            ss.State('ti_dead', float, np.nan),
+            ss.FloatArr('ti_exposed'),
+            ss.FloatArr('ti_symptomatic'),
+            ss.FloatArr('ti_recovered'),
+            ss.FloatArr('ti_dead'),
         )
 
         return
@@ -95,24 +95,23 @@ def update_pre(self, sim):
         Adapted from https://github.com/optimamodel/gavi-outbreaks/blob/main/stisim/gavi/cholera.py
         Original version by Dom Delport
         """
-
         # Progress exposed -> infected
-        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
+        infected = (self.exposed & (self.ti_infected <= sim.ti)).uids
         self.infected[infected] = True
 
         # Progress infected -> symptomatic
-        symptomatic = ss.true(self.infected & (self.ti_symptomatic <= sim.ti))
+        symptomatic = (self.infected & (self.ti_symptomatic <= sim.ti)).uids
         self.symptomatic[symptomatic] = True
 
         # Progress symptomatic -> recovered
-        recovered = ss.true(self.infectious & (self.ti_recovered <= sim.ti))
+        recovered = (self.infectious & (self.ti_recovered <= sim.ti)).uids
+        self.exposed[recovered] = False
         self.infected[recovered] = False
-        self.infectious[recovered] = False
         self.symptomatic[recovered] = False
         self.recovered[recovered] = True
 
         # Trigger deaths
-        deaths = ss.true(self.ti_dead <= sim.ti)
+        deaths = (self.ti_dead <= sim.ti).uids
         if len(deaths):
             sim.people.request_death(deaths)
 
@@ -128,8 +127,8 @@ def calc_environmental_prev(self, sim):
         p = self.pars
         r = self.results
 
-        n_symptomatic = len(ss.true(self.symptomatic))
-        n_asymptomatic = len(ss.true(self.asymptomatic))
+        n_symptomatic = self.symptomatic.sum()
+        n_asymptomatic = self.asymptomatic.sum()
         old_prev = self.results.env_prev[sim.ti-1]
 
         new_bacteria = p.shedding_rate * (n_symptomatic + p.asymp_trans * n_asymptomatic)
diff --git a/starsim/diseases/ebola.py b/starsim/diseases/ebola.py
index 1cdbadd3..162cf701 100644
--- a/starsim/diseases/ebola.py
+++ b/starsim/diseases/ebola.py
@@ -16,7 +16,7 @@ class Ebola(SIR):
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         """ Initialize with parameters """
 
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             # Natural history parameters, all specified in days
             dur_exp2symp = 12.7, # Add source
             dur_symp2sev =    6, # Add source
@@ -35,7 +35,7 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
             unburied_factor = 2.1,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_exp2symp    = ss.lognorm_ex,
             dur_symp2sev    = ss.lognorm_ex,
             dur_sev2dead    = ss.lognorm_ex,
@@ -54,17 +54,17 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         
         self.add_states(
             # SIR are added automatically, here we add E
-            ss.State('exposed', bool, False),
-            ss.State('severe', bool, False),
-            ss.State('recovered', bool, False),
-            ss.State('buried', bool, False),
+            ss.BoolArr('exposed'),
+            ss.BoolArr('severe'),
+            ss.BoolArr('recovered'),
+            ss.BoolArr('buried'),
     
             # Timepoint states
-            ss.State('ti_exposed', float, np.nan),
-            ss.State('ti_severe', float, np.nan),
-            ss.State('ti_recovered', float, np.nan),
-            ss.State('ti_dead', float, np.nan),
-            ss.State('ti_buried', float, np.nan),
+            ss.FloatArr('ti_exposed'),
+            ss.FloatArr('ti_severe'),
+            ss.FloatArr('ti_recovered'),
+            ss.FloatArr('ti_dead'),
+            ss.FloatArr('ti_buried'),
         )
 
         return
@@ -76,31 +76,31 @@ def infectious(self):
     def update_pre(self, sim):
 
         # Progress exposed -> infected
-        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
+        infected = (self.exposed & (self.ti_infected <= sim.ti)).uids
         self.exposed[infected] = False
         self.infected[infected] = True
 
         # Progress infectious -> severe
-        severe = ss.true(self.infected & (self.ti_severe <= sim.ti))
+        severe = (self.infected & (self.ti_severe <= sim.ti)).uids
         self.severe[severe] = True
 
         # Progress infected -> recovered
-        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        recovered = (self.infected & (self.ti_recovered <= sim.ti)).uids
         self.infected[recovered] = False
         self.recovered[recovered] = True
 
         # Progress severe -> recovered
-        recovered_sev = ss.true(self.severe & (self.ti_recovered <= sim.ti))
+        recovered_sev = (self.severe & (self.ti_recovered <= sim.ti)).uids
         self.severe[recovered_sev] = False
         self.recovered[recovered_sev] = True
 
         # Trigger deaths
-        deaths = ss.true(self.ti_dead <= sim.ti)
+        deaths = (self.ti_dead <= sim.ti).uids
         if len(deaths):
             sim.people.request_death(deaths)
 
         # Progress dead -> buried
-        buried = ss.true(self.ti_buried <= sim.ti)
+        buried = (self.ti_buried <= sim.ti).uids
         self.buried[buried] = True
         
         return
@@ -140,7 +140,7 @@ def set_prognoses(self, sim, uids, source_uids=None):
         # Change rel_trans values
         self.rel_trans[self.infectious] = 1
         self.rel_trans[self.severe] = self.pars['sev_factor']  # Change for severe
-        unburied_uids = ss.true((self.ti_dead <= sim.ti) & (self.ti_buried > sim.ti))
+        unburied_uids = ((self.ti_dead <= sim.ti) & (self.ti_buried > sim.ti)).uids
         self.rel_trans[unburied_uids] = self.pars['unburied_factor']  # Change for unburied
 
         return
diff --git a/starsim/diseases/gonorrhea.py b/starsim/diseases/gonorrhea.py
index 066fc946..6f6527f4 100644
--- a/starsim/diseases/gonorrhea.py
+++ b/starsim/diseases/gonorrhea.py
@@ -16,20 +16,20 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         # Additional states dependent on parameter values, e.g. self.p_symp?
         # These might be useful for connectors to target, e.g. if HIV reduces p_clear
         self.add_states(
-            ss.State('symptomatic', bool, False),
-            ss.State('ti_clearance', int, ss.INT_NAN),
-            ss.State('p_symp', float, 1),
+            ss.BoolArr('symptomatic'),
+            ss.FloatArr('ti_clearance'),
+            ss.FloatArr('p_symp', default=1),
         )
 
         # Parameters
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             dur_inf_in_days = ss.lognorm_ex(mean=10, stdev=0.6),  # median of 10 days (IQR 7–15 days) https://sti.bmj.com/content/96/8/556
             p_symp = 0.5,  # Share of infections that are symptomatic. Placeholder value
             p_clear = 0.2,  # Share of infections that spontaneously clear: https://sti.bmj.com/content/96/8/556
             init_prev = 0.1,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_inf_in_days = ss.lognorm_ex,
             p_symp          = ss.bernoulli,
             p_clear         = ss.bernoulli,
diff --git a/starsim/diseases/hiv.py b/starsim/diseases/hiv.py
index 181ea011..ef153287 100644
--- a/starsim/diseases/hiv.py
+++ b/starsim/diseases/hiv.py
@@ -15,13 +15,13 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
 
         # States
         self.add_states(
-            ss.State('on_art', bool, False),
-            ss.State('ti_art', int, ss.INT_NAN),
-            ss.State('cd4', float, 500),
-            ss.State('ti_dead', int, ss.INT_NAN), # Time of HIV-cause death
+            ss.BoolArr('on_art'),
+            ss.FloatArr('ti_art'),
+            ss.FloatArr('ti_dead'), # Time of HIV-cause death
+            ss.FloatArr('cd4', default=500),
         )
 
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             cd4_min = 100,
             cd4_max = 500,
             cd4_rate = 5,
@@ -31,7 +31,7 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
             death_prob = 0.05,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             init_prev  = ss.bernoulli,
             death_prob = ss.bernoulli,
         )
@@ -56,7 +56,7 @@ def update_pre(self, sim):
 
         self.rel_trans[sim.people.alive & self.infected & self.on_art] = 1 - self.pars['art_efficacy']
 
-        can_die = ss.true(sim.people.alive & sim.people.hiv.infected)
+        can_die = sim.people.hiv.infected.uids
         hiv_deaths = self.pars.death_prob.filter(can_die)
         
         sim.people.request_death(hiv_deaths)
diff --git a/starsim/diseases/measles.py b/starsim/diseases/measles.py
index 021179b8..bf7410f7 100644
--- a/starsim/diseases/measles.py
+++ b/starsim/diseases/measles.py
@@ -4,7 +4,6 @@
 Original version by @alina-muellenmeister, @domdelport, and @RomeshA
 """
 
-import numpy as np
 import starsim as ss
 from starsim.diseases.sir import SIR
 
@@ -16,7 +15,7 @@ class Measles(SIR):
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         """ Initialize with parameters """
 
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             # Natural history parameters, all specified in days
             dur_exp = 8,       # (days) - source: US CDC
             dur_inf = 11,      # (days) - source: US CDC
@@ -27,7 +26,7 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
             beta = None,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_exp   = ss.normal,
             dur_inf   = ss.normal,
             init_prev = ss.bernoulli,
@@ -38,8 +37,8 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
 
         # SIR are added automatically, here we add E
         self.add_states(
-            ss.State('exposed', bool, False),
-            ss.State('ti_exposed', float, np.nan),
+            ss.BoolArr('exposed'),
+            ss.FloatArr('ti_exposed'),
         )
 
         return
@@ -50,17 +49,17 @@ def infectious(self):
 
     def update_pre(self, sim):
         # Progress exposed -> infected
-        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
+        infected = (self.exposed & (self.ti_infected <= sim.ti)).uids
         self.exposed[infected] = False
         self.infected[infected] = True
 
         # Progress infected -> recovered
-        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        recovered = (self.infected & (self.ti_recovered <= sim.ti)).uids
         self.infected[recovered] = False
         self.recovered[recovered] = True
 
         # Trigger deaths
-        deaths = ss.true(self.ti_dead <= sim.ti)
+        deaths = (self.ti_dead <= sim.ti).uids
         if len(deaths):
             sim.people.request_death(deaths)
         return
diff --git a/starsim/diseases/ncd.py b/starsim/diseases/ncd.py
index dfa5fe75..e57d804b 100644
--- a/starsim/diseases/ncd.py
+++ b/starsim/diseases/ncd.py
@@ -25,12 +25,12 @@ def __init__(self, pars=None):
             prognosis = ss.weibull(c=2, scale=5), # Time in years between first becoming affected and death
         )
 
-        super().__init__(ss.omerge(default_pars, pars))
+        super().__init__(ss.dictmerge(default_pars, pars))
         self.add_states(
-            ss.State('at_risk', bool, False),
-            ss.State('affected', bool, False),
-            ss.State('ti_affected', int, ss.INT_NAN),
-            ss.State('ti_dead', int, ss.INT_NAN),
+            ss.BoolArr('at_risk'),
+            ss.BoolArr('affected'),
+            ss.FloatArr('ti_affected'),
+            ss.FloatArr('ti_dead'),
         )
         return
 
@@ -45,21 +45,20 @@ def set_initial_states(self, sim):
         i.e., creating their dynamic array, linking them to a People instance. That should have already
         taken place by the time this method is called.
         """
-        alive_uids = ss.true(sim.people.alive)
-        initial_risk = self.pars['initial_risk'].filter(alive_uids)
+        initial_risk = self.pars['initial_risk'].filter()
         self.at_risk[initial_risk] = True
         self.ti_affected[initial_risk] = sim.ti + sc.randround(self.pars['dur_risk'].rvs(initial_risk) / sim.dt)
         return initial_risk
 
     def update_pre(self, sim):
-        deaths = ss.true(self.ti_dead == sim.ti)
+        deaths = (self.ti_dead == sim.ti).uids
         sim.people.request_death(deaths)
         self.log.add_data(deaths, died=True)
         self.results.new_deaths[sim.ti] = len(deaths) # Log deaths attributable to this module
         return
 
     def make_new_cases(self, sim):
-        new_cases = ss.true(self.ti_affected == sim.ti)
+        new_cases = (self.ti_affected == sim.ti).uids
         self.affected[new_cases] = True
         prog_years = self.pars.prognosis.rvs(new_cases)
         self.ti_dead[new_cases] = sim.ti + sc.randround(prog_years / sim.dt)
@@ -81,8 +80,7 @@ def init_results(self, sim):
     def update_results(self, sim):
         super().update_results(sim)
         ti = sim.ti
-        alive = sim.people.alive
-        self.results.n_not_at_risk[ti] = np.count_nonzero(self.not_at_risk & alive)
-        self.results.prevalence[ti]    = np.count_nonzero(self.affected & alive)/alive.count()
+        self.results.n_not_at_risk[ti] = np.count_nonzero(self.not_at_risk)
+        self.results.prevalence[ti]    = np.count_nonzero(self.affected)/len(sim.people)
         self.results.new_deaths[ti]    = np.count_nonzero(self.ti_dead == ti)
         return
diff --git a/starsim/diseases/sir.py b/starsim/diseases/sir.py
index 2de3f2a8..cd8f7f4d 100644
--- a/starsim/diseases/sir.py
+++ b/starsim/diseases/sir.py
@@ -19,14 +19,14 @@ class SIR(ss.Infection):
     """
 
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             dur_inf = 6,
             init_prev = 0.01,
             p_death = 0.01,
             beta = 0.5,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_inf   = ss.lognorm_ex,
             init_prev = ss.bernoulli,
             p_death   = ss.bernoulli,
@@ -35,20 +35,20 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
         super().__init__(pars=pars, par_dists=par_dists, *args, **kwargs)
 
         self.add_states(
-            ss.State('recovered', bool, False),
-            ss.State('ti_recovered', int, ss.INT_NAN),
-            ss.State('ti_dead', int, ss.INT_NAN),
+            ss.BoolArr('recovered'),
+            ss.FloatArr('ti_recovered'),
+            ss.FloatArr('ti_dead'),
         )
         return
 
     def update_pre(self, sim):
         # Progress infectious -> recovered
-        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        recovered = (self.infected & (self.ti_recovered <= sim.ti)).uids
         self.infected[recovered] = False
         self.recovered[recovered] = True
 
         # Trigger deaths
-        deaths = ss.true(self.ti_dead <= sim.ti)
+        deaths = (self.ti_dead <= sim.ti).uids
         if len(deaths):
             sim.people.request_death(deaths)
         return
@@ -99,7 +99,7 @@ class SIS(ss.Infection):
     is no death in this case.
     """
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             dur_inf = 10,
             init_prev = 0.01,
             beta = 0.05,
@@ -107,14 +107,14 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
             imm_boost = 1.0,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             dur_inf   = ss.lognorm_ex,
             init_prev = ss.bernoulli,
         )
         
         self.add_states(
-            ss.State('ti_recovered', int, ss.INT_NAN),
-            ss.State('immunity', float, 0.0),
+            ss.FloatArr('ti_recovered'),
+            ss.FloatArr('immunity', default=0.0),
         )
 
         super().__init__(pars=pars, par_dists=par_dists, *args, **kwargs)
@@ -122,16 +122,16 @@ def __init__(self, pars=None, par_dists=None, *args, **kwargs):
 
     def update_pre(self, sim):
         """ Progress infectious -> recovered """
-        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        recovered = (self.infected & (self.ti_recovered <= sim.ti)).uids
         self.infected[recovered] = False
         self.susceptible[recovered] = True
         self.update_immunity(sim)
         return
     
     def update_immunity(self, sim):
-        uids = ss.true(self.immunity > 0)
-        self.immunity[uids] = (self.immunity[uids])*(1 - self.pars.waning*sim.dt)
-        self.rel_sus[uids] = np.maximum(0, 1 - self.immunity[uids])
+        has_imm = (self.immunity > 0).uids
+        self.immunity[has_imm] = (self.immunity[has_imm])*(1 - self.pars.waning*sim.dt)
+        self.rel_sus[has_imm] = np.maximum(0, 1 - self.immunity[has_imm])
         return
 
     def set_prognoses(self, sim, uids, source_uids=None):
@@ -179,7 +179,7 @@ class sir_vaccine(ss.Vx):
     Create a vaccine product that changes susceptible people to recovered (i.e., perfect immunity)
     """
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
-        pars = ss.omerge({
+        pars = ss.dictmerge({
             'efficacy': 0.9,
         }, pars)
 
diff --git a/starsim/diseases/syphilis.py b/starsim/diseases/syphilis.py
index f74416c8..abc95196 100644
--- a/starsim/diseases/syphilis.py
+++ b/starsim/diseases/syphilis.py
@@ -17,30 +17,28 @@ def __init__(self, pars=None):
 
         self.add_states(
             # Adult syphilis states
-            ss.State('exposed', bool, False),  # AKA incubating. Free of symptoms, not transmissible
-            ss.State('primary', bool, False),  # Primary chancres
-            ss.State('secondary', bool, False),  # Inclusive of those who may still have primary chancres
-            ss.State('latent_temp', bool, False),  # Relapses to secondary (~1y)
-            ss.State('latent_long', bool, False),  # Can progress to tertiary or remain here
-            ss.State('tertiary', bool, False),  # Includes complications (cardio/neuro/disfigurement)
-            ss.State('immune', bool, False),  # After effective treatment people may acquire temp immunity
-            ss.State('ever_exposed', bool, False),  # Anyone ever exposed - stays true after treatment
-    
-            # Congenital syphilis states
-            ss.State('congenital', bool, False),
+            ss.BoolArr('exposed'),  # AKA incubating. Free of symptoms, not transmissible
+            ss.BoolArr('primary'),  # Primary chancres
+            ss.BoolArr('secondary'),  # Inclusive of those who may still have primary chancres
+            ss.BoolArr('latent_temp'),  # Relapses to secondary (~1y)
+            ss.BoolArr('latent_long'),  # Can progress to tertiary or remain here
+            ss.BoolArr('tertiary'),  # Includes complications (cardio/neuro/disfigurement)
+            ss.BoolArr('immune'),  # After effective treatment people may acquire temp immunity
+            ss.BoolArr('ever_exposed'),  # Anyone ever exposed - stays true after treatment
+            ss.BoolArr('congenital'),  # Congenital syphilis states
     
             # Timestep of state changes
-            ss.State('ti_exposed', int, ss.INT_NAN),
-            ss.State('ti_primary', int, ss.INT_NAN),
-            ss.State('ti_secondary', int, ss.INT_NAN),
-            ss.State('ti_latent_temp', int, ss.INT_NAN),
-            ss.State('ti_latent_long', int, ss.INT_NAN),
-            ss.State('ti_tertiary', int, ss.INT_NAN),
-            ss.State('ti_immune', int, ss.INT_NAN),
-            ss.State('ti_miscarriage', int, ss.INT_NAN),
-            ss.State('ti_nnd', int, ss.INT_NAN),
-            ss.State('ti_stillborn', int, ss.INT_NAN),
-            ss.State('ti_congenital', int, ss.INT_NAN),
+            ss.FloatArr('ti_exposed'),
+            ss.FloatArr('ti_primary'),
+            ss.FloatArr('ti_secondary'),
+            ss.FloatArr('ti_latent_temp'),
+            ss.FloatArr('ti_latent_long'),
+            ss.FloatArr('ti_tertiary'),
+            ss.FloatArr('ti_immune'),
+            ss.FloatArr('ti_miscarriage'),
+            ss.FloatArr('ti_nnd'),
+            ss.FloatArr('ti_stillborn'),
+            ss.FloatArr('ti_congenital'),
         )
 
         # Parameters
@@ -80,7 +78,7 @@ def __init__(self, pars=None):
             # Initial conditions
             init_prev=ss.bernoulli(p=0.03),
         )
-        self.pars = ss.omerge(default_pars, self.pars) # NB: regular omerge rather than omergeleft
+        self.pars = ss.dictmerge(default_pars, self.pars) # NB: regular dictmerge rather than dictmergeleft
 
         return
 
@@ -129,11 +127,11 @@ def update_pre(self, sim):
         self.rel_trans[primary] = self.pars.rel_trans['primary']
 
         # Secondary from primary
-        secondary_from_primary = self.primary & (self.ti_secondary <= sim.ti)
-        if len(ss.true(secondary_from_primary)) > 0:
+        secondary_from_primary = (self.primary & (self.ti_secondary <= sim.ti)).uids
+        if len(secondary_from_primary) > 0:
             self.secondary[secondary_from_primary] = True
             self.primary[secondary_from_primary] = False
-            self.set_secondary_prognoses(sim, ss.true(secondary_from_primary))
+            self.set_secondary_prognoses(sim, secondary_from_primary)
             self.rel_trans[secondary_from_primary] = self.pars.rel_trans['secondary']
 
         # Hack to reset the MultiRNGs in set_secondary_prognoses so that they can be called again in this timestep. TODO: Refactor
@@ -141,38 +139,38 @@ def update_pre(self, sim):
         self.pars.dur_secondary.jump(sim.ti+1)
 
         # Secondary reactivation from latent
-        secondary_from_latent = self.latent_temp & (self.ti_secondary <= sim.ti)
-        if len(ss.true(secondary_from_latent)) > 0:
+        secondary_from_latent = (self.latent_temp & (self.ti_secondary <= sim.ti)).uids
+        if len(secondary_from_latent) > 0:
             self.secondary[secondary_from_latent] = True
             self.latent_temp[secondary_from_latent] = False
-            self.set_secondary_prognoses(sim, ss.true(secondary_from_latent))
+            self.set_secondary_prognoses(sim, secondary_from_latent)
             self.rel_trans[secondary_from_latent] = self.pars.rel_trans['secondary']
 
         # Latent
-        latent_temp = self.secondary & (self.ti_latent_temp <= sim.ti)
-        if len(ss.true(latent_temp)) > 0:
+        latent_temp = (self.secondary & (self.ti_latent_temp <= sim.ti)).uids
+        if len(latent_temp) > 0:
             self.latent_temp[latent_temp] = True
             self.secondary[latent_temp] = False
-            self.set_latent_temp_prognoses(sim, ss.true(latent_temp))
+            self.set_latent_temp_prognoses(sim, latent_temp)
             self.rel_trans[latent_temp] = self.pars.rel_trans['latent_temp']
 
         # Latent long
-        latent_long = self.secondary & (self.ti_latent_long <= sim.ti)
-        if len(ss.true(latent_long)) > 0:
+        latent_long = (self.secondary & (self.ti_latent_long <= sim.ti)).uids
+        if len(latent_long) > 0:
             self.latent_long[latent_long] = True
             self.secondary[latent_long] = False
-            self.set_latent_long_prognoses(sim, ss.true(latent_long))
+            self.set_latent_long_prognoses(sim, latent_long)
             self.rel_trans[latent_long] = self.pars.rel_trans['latent_long']
 
         # Tertiary
-        tertiary = self.latent_long & (self.ti_tertiary <= sim.ti)
+        tertiary = (self.latent_long & (self.ti_tertiary <= sim.ti)).uids
         self.tertiary[tertiary] = True
         self.latent_long[tertiary] = False
         self.rel_trans[tertiary] = self.pars.rel_trans['tertiary']
 
         # Congenital syphilis deaths
-        nnd = self.ti_nnd == sim.ti
-        stillborn = self.ti_stillborn == sim.ti
+        nnd = (self.ti_nnd == sim.ti).uids
+        stillborn = (self.ti_stillborn == sim.ti).uids
         sim.people.request_death(nnd)
         sim.people.request_death(stillborn)
 
@@ -265,7 +263,7 @@ def set_congenital(self, sim, target_uids, source_uids=None):
         # Determine outcomes
         for state in ['active', 'latent']:
 
-            source_state_inds = getattr(self, state)[source_uids].values.nonzero()[-1]
+            source_state_inds = getattr(self, state)[source_uids].nonzero()[0]
             uids = target_uids[source_state_inds]
 
             if len(uids) > 0:
@@ -273,7 +271,7 @@ def set_congenital(self, sim, target_uids, source_uids=None):
                 # Birth outcomes must be modified to add probability of susceptible birth
                 birth_outcomes = self.pars.birth_outcomes[state]
                 assigned_outcomes = birth_outcomes.rvs(len(uids))
-                time_to_birth = -sim.people.age
+                time_to_birth = -sim.people.age.raw # TODO: make nicer
 
                 # Schedule events
                 for oi, outcome in enumerate(self.pars.birth_outcome_keys):
@@ -281,7 +279,7 @@ def set_congenital(self, sim, target_uids, source_uids=None):
                     if len(o_uids) > 0:
                         ti_outcome = f'ti_{outcome}'
                         vals = getattr(self, ti_outcome)
-                        vals[o_uids] = sim.ti + rr(time_to_birth[o_uids].values / sim.dt)
+                        vals[o_uids] = sim.ti + rr(time_to_birth[o_uids] / sim.dt)
                         setattr(self, ti_outcome, vals)
 
         return
@@ -343,7 +341,7 @@ def check_eligibility(self, sim):
         if self.eligibility is not None:
             is_eligible = self.eligibility(sim)
         else:
-            is_eligible = sim.people.alive  # Probably not required
+            is_eligible = sim.people.auids  # Probably not required
         return is_eligible
 
     def initialize(self, sim):
diff --git a/starsim/distributions.py b/starsim/distributions.py
index b3793516..e4f1d31c 100644
--- a/starsim/distributions.py
+++ b/starsim/distributions.py
@@ -370,7 +370,7 @@ def process_size(self, n=1):
             slots = None
             size = n
         else:
-            uids = np.asarray(n)
+            uids = ss.uids(n)
             if len(uids):
                 slots = self.slots[uids]
                 if len(slots): # Handle case where uids is boolean
@@ -698,8 +698,10 @@ def ppf(self, rands):
         rvs = rands < self._pars.p
         return rvs
     
-    def filter(self, uids, both=False):
+    def filter(self, uids=None, both=False):
         """ Return UIDs that correspond to True, or optionally return both True and False """
+        if uids is None:
+            uids = self.sim.people.auids # All active UIDs
         bools = self.rvs(uids)
         if both:
             return uids[bools], uids[~bools]
diff --git a/starsim/interventions.py b/starsim/interventions.py
index 90ac23df..df8ca413 100644
--- a/starsim/interventions.py
+++ b/starsim/interventions.py
@@ -81,8 +81,14 @@ def check_eligibility(self, sim):
         """
         if self.eligibility is not None:
             is_eligible = self.eligibility(sim)
+            if is_eligible is not None and len(is_eligible): # Only worry if non-None/nonzero length
+                if isinstance(is_eligible, ss.BoolArr):
+                    is_eligible = is_eligible.uids
+                if not isinstance(is_eligible, ss.uids):
+                    errormsg = f'Eligibility function must return BoolArr or UIDs, not {type(is_eligible)} {is_eligible}'
+                    raise TypeError(errormsg)
         else:
-            is_eligible = sim.people.alive  # Probably not required
+            is_eligible = sim.people.auids # Everyone
         return is_eligible
 
 
@@ -198,9 +204,9 @@ def __init__(self, product=None, prob=None, eligibility=None, **kwargs):
         self.prob = sc.promotetoarray(prob)
         self.eligibility = eligibility
         self._parse_product(product)
-        self.screened = ss.State('screened', bool, False)
-        self.screens = ss.State('screens', int, 0)
-        self.ti_screened = ss.State('ti_screened', int, ss.INT_NAN)
+        self.screened = ss.BoolArr('screened')
+        self.screens = ss.FloatArr('screens', default=0)
+        self.ti_screened = ss.FloatArr('ti_screened')
         return
 
     def initialize(self, sim):
@@ -214,9 +220,9 @@ def deliver(self, sim):
         """
         ti = sc.findinds(self.timepoints, sim.ti)[0]
         prob = self.prob[ti]  # Get the proportion of people who will be tested this timestep
-        is_eligible = self.check_eligibility(sim)  # Check eligibility
+        eligible_uids = self.check_eligibility(sim)  # Check eligibility
         self.coverage_dist.set(p=prob)
-        accept_uids = self.coverage_dist.filter(ss.true(is_eligible))
+        accept_uids = self.coverage_dist.filter(eligible_uids)
         if len(accept_uids):
             self.outcomes = self.product.administer(sim, accept_uids)  # Actually administer the diagnostic
         return accept_uids
@@ -241,7 +247,7 @@ def apply(self, sim, module=None):
         """
         Perform screening by finding who's eligible, finding who accepts, and applying the product.
         """
-        accept_uids = np.array([])
+        accept_uids = ss.uids()
         if sim.ti in self.timepoints:
             accept_uids = self.deliver(sim)
             self.screened[accept_uids] = True
@@ -264,7 +270,7 @@ def check_eligibility(self, sim):
 
     def apply(self, sim):
         self.outcomes = {k: np.array([], dtype=int) for k in self.product.hierarchy}
-        accept_inds = np.array([])
+        accept_inds = ss.uids()
         if sim.t in self.timepoints: accept_inds = self.deliver(sim)
         return accept_inds
 
@@ -393,7 +399,7 @@ def get_accept_inds(self, sim):
         """
         Get indices of people who will acccept treatment; these people are then added to a queue or scheduled for receiving treatment
         """
-        accept_uids = np.array([], dtype=int)
+        accept_uids = ss.uids()
         eligible_uids = self.check_eligibility(sim)  # Apply eligiblity
         if len(eligible_uids):
             self.coverage_dist.set(p=self.prob[0])
@@ -413,7 +419,7 @@ def apply(self, sim):
         # Get indices of who will get treated
         treat_candidates = self.get_candidates(sim)  # NB, this needs to be implemented by derived classes
         still_eligible = self.check_eligibility(sim)
-        treat_uids = np.intersect1d(treat_candidates, still_eligible)
+        treat_uids = treat_candidates.intersect(still_eligible)
         if len(treat_uids):
             self.outcomes = self.product.administer(sim, treat_uids)
         return treat_uids
@@ -450,7 +456,7 @@ def get_candidates(self, sim):
                 treat_candidates = self.queue[:]
             else:
                 treat_candidates = self.queue[:self.max_capacity]
-        return sc.promotetoarray(treat_candidates)
+        return ss.uids(treat_candidates) # TODO: Check
 
     def apply(self, sim):
         """
@@ -483,9 +489,9 @@ def __init__(self, product=None, prob=None, label=None, **kwargs):
         self.prob = sc.promotetoarray(prob)
         self.label = label
         self._parse_product(product)
-        self.vaccinated = ss.State('vaccinated', bool, False)
-        self.n_doses = ss.State('doses', int, 0)
-        self.ti_vaccinated = ss.State('ti_vaccinated', int, ss.INT_NAN)
+        self.vaccinated = ss.BoolArr('vaccinated')
+        self.n_doses = ss.FloatArr('doses', default=0)
+        self.ti_vaccinated = ss.FloatArr('ti_vaccinated')
         self.coverage_dist = ss.bernoulli(p=0)  # Placeholder
         return
 
@@ -500,7 +506,7 @@ def apply(self, sim):
             prob = self.prob[ti]  # Get the proportion of people who will be tested this timestep
             is_eligible = self.check_eligibility(sim)  # Check eligibility
             self.coverage_dist.set(p=prob)
-            accept_uids = self.coverage_dist.filter(ss.true(is_eligible))
+            accept_uids = self.coverage_dist.filter(is_eligible)
 
             if len(accept_uids):
                 self.product.administer(sim.people, accept_uids)
diff --git a/starsim/modules.py b/starsim/modules.py
index 22a51864..20218986 100644
--- a/starsim/modules.py
+++ b/starsim/modules.py
@@ -12,8 +12,8 @@
 class Module(sc.quickobj):
 
     def __init__(self, pars=None, par_dists=None, name=None, label=None, requires=None, **kwargs):
-        self.pars = ss.omerge(pars, kwargs)
-        self.par_dists = ss.omerge(par_dists)
+        self.pars = ss.dictmerge(pars, kwargs)
+        self.par_dists = ss.dictmerge(par_dists)
         self.name = sc.ifelse(name, getattr(self, 'name', self.__class__.__name__.lower())) # Default name is the class name
         self.label = sc.ifelse(label, getattr(self, 'label', self.name))
         self.requires = sc.mergelists(requires)
@@ -142,7 +142,7 @@ def add_states(self, *args, check=True):
                 state = arg
                 
             if check:
-                assert isinstance(state, ss.State), f'Could not add {state}: not a State object'
+                assert isinstance(state, ss.Arr), f'Could not add {state}: not an Arr object'
                 
             setattr(self, state.name, state)
         return
@@ -158,7 +158,7 @@ def states(self):
         due to supporting features like multiple genotypes) then the Module should
         overload this attribute to ensure that all states appear in here.
         """
-        return [x for x in self.__dict__.values() if isinstance(x, ss.State)]
+        return [x for x in self.__dict__.values() if isinstance(x, ss.Arr)]
 
     @property
     def statesdict(self):
diff --git a/starsim/network.py b/starsim/network.py
index 3351925a..018361b1 100644
--- a/starsim/network.py
+++ b/starsim/network.py
@@ -71,7 +71,7 @@ def __init__(self, pars=None, key_dict=None, vertical=False, *args, **kwargs):
             p2 = ss_int_,
             beta = ss_float_,
         )
-        self.meta = ss.omerge(default_keys, key_dict)
+        self.meta = ss.dictmerge(default_keys, key_dict)
         self.vertical = vertical  # Whether transmission is bidirectional
 
         # Initialize the keys of the network
@@ -85,7 +85,8 @@ def __init__(self, pars=None, key_dict=None, vertical=False, *args, **kwargs):
             self.initialized = True
 
         # Define states using placeholder values
-        self.participant = ss.State('participant', bool, default=False)
+        self.participant = ss.BoolArr('participant')
+        self.validate_uids()
         return
     
     @property
@@ -128,12 +129,12 @@ def __contains__(self, item):
 
         Returns: True if person index appears in any interactions
         """
-        return (item in self.contacts.p1) or (item in self.contacts.p2)
+        return (item in self.contacts.p1) or (item in self.contacts.p2) # TODO: chek if (item in self.members) is faster
 
     @property
     def members(self):
         """ Return sorted array of all members """
-        return np.unique([self.contacts.p1, self.contacts.p2])
+        return np.unique([self.contacts.p1, self.contacts.p2]).view(ss.uids)
 
     def meta_keys(self):
         """ Return the keys for the network's meta information """
@@ -147,6 +148,16 @@ def set_network_states(self, people):
         states that depend on other states.
         """
         pass
+    
+    def validate_uids(self):
+        """ Ensure that p1, p2 are both UID arrays """
+        contacts = self.contacts
+        for key in ['p1', 'p2']:
+            if key in contacts:
+                arr = contacts[key]
+                if not isinstance(arr, ss.uids):
+                    self.contacts[key] = ss.uids(arr)
+        return
 
     def validate(self, force=True):
         """
@@ -166,6 +177,7 @@ def validate(self, force=True):
             if n != actual_n:
                 errormsg = f'Expecting length {n} for network key "{key}"; got {actual_n}'  # Report length mismatches
                 raise TypeError(errormsg)
+        self.validate_uids()
         return
 
     def get_inds(self, inds, remove=False):
@@ -180,6 +192,7 @@ def get_inds(self, inds, remove=False):
             output[key] = self.contacts[key][inds]  # Copy to the output object
             if remove:
                 self.contacts[key] = np.delete(self.contacts[key], inds)  # Remove from the original
+                self.validate_uids()
         return output
 
     def pop_inds(self, inds):
@@ -193,20 +206,23 @@ def pop_inds(self, inds):
         popped_inds = self.get_inds(inds, remove=True)
         return popped_inds
 
-    def append(self, contacts):
+    def append(self, contacts=None, **kwargs):
         """
         Append contacts to the current network.
 
         Args:
             contacts (dict): a dictionary of arrays with keys p1,p2,beta, as returned from network.pop_inds()
         """
+        contacts = sc.mergedicts(contacts, kwargs)
         for key in self.meta_keys():
-            new_arr = contacts[key]
-            n_curr = len(self.contacts[key])  # Current number of contacts
-            n_new = len(new_arr)  # New contacts to add
-            n_total = n_curr + n_new  # New size
-            self.contacts[key] = np.resize(self.contacts[key], n_total)  # Resize to make room, preserving dtype
-            self.contacts[key][n_curr:] = new_arr  # Copy contacts into the network
+            curr_arr = self.contacts[key]
+            try:
+                new_arr = contacts[key]
+            except KeyError:
+                errormsg = f'Cannot append contacts since required key "{key}" is missing'
+                raise KeyError(errormsg)
+            self.contacts[key] = np.concatenate([curr_arr, new_arr])  # Resize to make room, preserving dtype
+        self.validate_uids()
         return
 
     def to_dict(self):
@@ -319,7 +335,7 @@ def update(self, people):
 class DynamicNetwork(Network):
     """ A network where partnerships update dynamically """
     def __init__(self, pars=None, key_dict=None, **kwargs):
-        key_dict = ss.omerge({'dur': ss_float_}, key_dict)
+        key_dict = ss.dictmerge({'dur': ss_float_}, key_dict)
         super().__init__(pars, key_dict=key_dict, **kwargs)
         return
 
@@ -337,14 +353,16 @@ def end_pairs(self, people):
 class SexualNetwork(Network):
     """ Base class for all sexual networks """
     def __init__(self, pars=None, key_dict=None, **kwargs):
-        key_dict = ss.omerge({'acts': ss_int_}, key_dict)
+        key_dict = ss.dictmerge({'acts': ss_int_}, key_dict)
         super().__init__(pars, key_dict=key_dict, **kwargs)
-        self.debut = ss.State('debut', float, default=0)
+        self.debut = ss.FloatArr('debut', default=0)
         return
 
     def active(self, people):
         # Exclude people who are not alive
-        return self.participant & (people.age > self.debut) & people.alive
+        valid_age = people.age > self.debut
+        active = self.participant & valid_age & people.alive
+        return active
 
     def available(self, people, sex):
         # Currently assumes unpartnered people are available
@@ -352,7 +370,11 @@ def available(self, people, sex):
         # This property could also be overwritten by a NetworkConnector
         # which could incorporate information about membership in other
         # contact networks
-        return np.setdiff1d(ss.true(people[sex] & self.active(people)), self.members) # ss.true instead of people.uid[]?
+        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
 
     def beta_per_dt(self, disease_beta=None, dt=None, uids=None):
         if uids is None: uids = Ellipsis
@@ -388,7 +410,7 @@ class StaticNet(Network):
     def __init__(self, graph=None, pars=None, **kwargs):
         super().__init__(**kwargs)
         self.graph = graph
-        self.pars = ss.omerge(dict(seed=True), pars)
+        self.pars = ss.dictmerge(dict(seed=True), pars)
         self.dist = ss.Dist(name='StaticNet').initialize()
         return
 
@@ -431,9 +453,8 @@ def get_contacts(self):
             p1, p2 = edge
             p1s.append(p1)
             p2s.append(p2)
-        self.contacts.p1 = np.concatenate([self.contacts.p1, p1s])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, p2s])
-        self.contacts.beta = np.concatenate([self.contacts.beta, np.ones_like(p1s)])
+        contacts = dict(p1=p1s, p2=p2s, beta=np.ones_like(p1s))
+        self.append(contacts)
         return
 
 
@@ -442,7 +463,7 @@ class RandomNet(DynamicNetwork):
 
     def __init__(self, pars=None, par_dists=None, key_dict=None, **kwargs):
         """ Initialize """
-        pars = ss.omerge({
+        pars = ss.dictmerge({
             'n_contacts': 10,  # Distribution or int. If int, interpreted as the mean of the dist listed in par_dists
             'dur': 0,
         }, pars)
@@ -520,12 +541,8 @@ def add_pairs(self, people):
             dur = self.pars.dur.rvs(p1)
         else:
             dur = np.full(len(p1), self.pars.dur)
-
-        self.contacts.p1 = np.concatenate([self.contacts.p1, p1])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, p2])
-        self.contacts.beta = np.concatenate([self.contacts.beta, beta])
-        self.contacts.dur = np.concatenate([self.contacts.dur, dur])
-
+        
+        self.append(p1=p1, p2=p2, beta=beta, dur=dur)
         return
 
 
@@ -564,9 +581,7 @@ def initialize(self, sim):
 
     def get_contacts(self):
         indices = np.arange(self.n)
-        self.contacts.p1 = np.concatenate([self.contacts.p1, indices])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, indices])
-        self.contacts.beta = np.concatenate([self.contacts.beta, np.zeros_like(indices)])
+        self.append(dict(p1=indices, p2=indices, beta=np.zeros_like(indices)))
         return
 
 
@@ -577,7 +592,7 @@ class MFNet(SexualNetwork, DynamicNetwork):
     """
 
     def __init__(self, pars=None, par_dists=None, key_dict=None, **kwargs):
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             duration = 15,  # Can vary by age, year, and individual pair. Set scale=exp(mu) and s=sigma where mu,sigma are of the underlying normal distribution.
             participation = 0.9,  # Probability of participating in this network - can vary by individual properties (age, sex, ...) using callable parameter values
             debut = 16,  # Age of debut can vary by using callable parameter values
@@ -585,7 +600,7 @@ def __init__(self, pars=None, par_dists=None, key_dict=None, **kwargs):
             rel_part_rates = 1.0,
         )
 
-        par_dists = ss.omergeleft(par_dists,
+        par_dists = ss.dictmergeleft(par_dists,
             duration      = ss.lognorm_ex,
             participation = ss.bernoulli,
             debut         = ss.normal,
@@ -614,15 +629,15 @@ def set_network_states(self, people, upper_age=None):
 
     def set_participation(self, people, upper_age=None):
         # Set people who will participate in the network at some point
-        if upper_age is None: uids = people.uid
-        else: uids = people.uid[(people.age < upper_age)]
+        if upper_age is None: uids = people.auids
+        else: uids = (people.age < upper_age).uids
         self.participant[uids] = self.pars.participation.rvs(uids)
         return
 
     def set_debut(self, people, upper_age=None):
         # Set debut age
-        if upper_age is None: uids = people.uid
-        else: uids = people.uid[(people.age < upper_age)]
+        if upper_age is None: uids = people.auids
+        else: uids = (people.age < upper_age).uids
         self.debut[uids] = self.pars.debut.rvs(uids)
         return
 
@@ -658,11 +673,7 @@ def add_pairs(self, people, ti=None):
             dur_vals = self.pars.duration.rvs(len(p1))  # Just use len(p1) to say how many draws are needed
             act_vals = self.pars.acts.rvs(len(p1))
 
-        self.contacts.p1 = np.concatenate([self.contacts.p1, p1])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, p2])
-        self.contacts.beta = np.concatenate([self.contacts.beta, beta])
-        self.contacts.dur = np.concatenate([self.contacts.dur, dur_vals])
-        self.contacts.acts = np.concatenate([self.contacts.acts, act_vals])
+        self.append(p1=p1, p2=p2, beta=beta, dur=dur_vals, acts=act_vals)
 
         return len(p1)
 
@@ -679,7 +690,7 @@ class MSMNet(SexualNetwork, DynamicNetwork):
     """
 
     def __init__(self, pars=None, key_dict=None, **kwargs):
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             duration_dist = ss.lognorm_ex(mean=15, stdev=15),
             participation_dist = ss.bernoulli(p=0.1),  # Probability of participating in this network - can vary by individual properties (age, sex, ...) using callable parameter values
             debut_dist = ss.normal(loc=16, scale=2),
@@ -731,11 +742,8 @@ def add_pairs(self, people, ti=None):
             dur = self.pars.duration.rvs(len(p1)) # Just use len(p1) to say how many draws are needed
             act_vals = self.pars.acts.rvs(len(p1))
 
-        self.contacts.p1 = np.concatenate([self.contacts.p1, p1])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, p2])
-        self.contacts.beta = np.concatenate([self.contacts.beta, np.ones_like(p1)])
-        self.contacts.dur = np.concatenate([self.contacts.dur, dur])
-        self.contacts.acts = np.concatenate([self.contacts.acts, act_vals])
+        self.append(p1=p1, p2=p2, beta=np.ones_like(p1), dur=dur, acts=act_vals)
+        
         return len(p1)
 
     def update(self, people, dt=None):
@@ -758,7 +766,7 @@ def __init__(self, pars=None, **kwargs):
         std is the standard deviation of noise added to the age of each individual when seeking a pair. Larger values will created more diversity in age gaps.
         
         """
-        pars = ss.omerge({
+        pars = ss.dictmerge({
             'embedding_func': ss.normal(name='EmbeddingNet', loc=self.embedding_loc, scale=2),
             'male_shift': 5,
         }, pars)
@@ -767,7 +775,7 @@ def __init__(self, pars=None, **kwargs):
 
     @staticmethod
     def embedding_loc(module, sim, uids):
-        loc = sim.people.age[uids].values
+        loc = sim.people.age[uids]
         loc[sim.people.female[uids]] += module.pars.male_shift  # Shift females so they will be paired with older men
         return loc
 
@@ -780,7 +788,7 @@ def add_pairs(self, people, ti=None):
                 print('No pairs to add')
             return 0
 
-        available = np.concatenate((available_m, available_f))
+        available = ss.uids.cat(available_m, available_f)
         loc = self.pars.embedding_func.rvs(available)
         loc_f = loc[people.female[available]]
         loc_m = loc[~people.female[available]]
@@ -788,22 +796,17 @@ def add_pairs(self, people, ti=None):
         dist_mat = spsp.distance_matrix(loc_m[:, np.newaxis], loc_f[:, np.newaxis])
 
         ind_m, ind_f = spo.linear_sum_assignment(dist_mat)
-        # loc_f[ind_f[0]] is close to loc_m[ind_m[0]]
 
         n_pairs = len(ind_f)
 
-        beta = np.ones(n_pairs)
-
-        # Figure out durations
+        # Finalize pairs
         p1 = available_m[ind_m]
+        p2 = available_f[ind_f]
+        beta = np.ones(n_pairs) # TODO: Allow custom beta
         dur_vals = self.pars.duration.rvs(p1)
         act_vals = self.pars.acts.rvs(p1)
 
-        self.contacts.p1 = np.concatenate([self.contacts.p1, p1])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, available_f[ind_f]])
-        self.contacts.beta = np.concatenate([self.contacts.beta, beta])
-        self.contacts.dur = np.concatenate([self.contacts.dur, dur_vals])
-        self.contacts.acts = np.concatenate([self.contacts.acts, act_vals])
+        self.append(p1=p1, p2=p2, beta=beta, dur=dur_vals, acts=act_vals)
         return len(beta)
 
 
@@ -836,12 +839,10 @@ def add_pairs(self, mother_inds, unborn_inds, dur):
         """
         Add connections between pregnant women and their as-yet-unborn babies
         """
-        beta = np.ones_like(mother_inds)
-        self.contacts.p1 = np.concatenate([self.contacts.p1, mother_inds])
-        self.contacts.p2 = np.concatenate([self.contacts.p2, unborn_inds])
-        self.contacts.beta = np.concatenate([self.contacts.beta, beta])
-        self.contacts.dur = np.concatenate([self.contacts.dur, dur])
-        return len(mother_inds)
+        n = len(mother_inds)
+        beta = np.ones(n)
+        self.append(p1=mother_inds, p2=unborn_inds, beta=beta, dur=dur)
+        return n
 
 
 # %% Network connectors
@@ -867,7 +868,7 @@ class MF_MSM(NetworkConnector):
     """ Combines the MF and MSM networks """
     def __init__(self, pars=None):
         networks = [ss.MFNet, ss.MSMNet]
-        pars = ss.omergeleft(pars,
+        pars = ss.dictmergeleft(pars,
             prop_bi = 0.5,  # Could vary over time -- but not by age or sex or individual
         )
         super().__init__(networks=networks, pars=pars)
diff --git a/starsim/parameters.py b/starsim/parameters.py
index 5eb01870..71b0d2fc 100644
--- a/starsim/parameters.py
+++ b/starsim/parameters.py
@@ -27,7 +27,6 @@ def __init__(self, **kwargs):
         self.n_agents        = 10e3  # Number of agents
         self.total_pop       = None  # If defined, used for calculating the scale factor
         self.pop_scale       = None  # How much to scale the population
-        self.remove_dead     = 10    # How many timesteps to go between removing dead agents (0 to not remove)
 
         # Demographic parameters
         self.location    = None  #  NOT CURRENTLY FUNCTIONAL - what demographics to use
diff --git a/starsim/people.py b/starsim/people.py
index 5cf1dcbd..3ab02290 100644
--- a/starsim/people.py
+++ b/starsim/people.py
@@ -22,27 +22,21 @@ class BasePeople(sc.prettyobj):
 
     def __init__(self, n_agents):
 
-        self.initialized = False
-        self._uid_map = ss.ArrayView(int, default=ss.INT_NAN)  # This variable tracks all UIDs ever created
-        self.uid = ss.ArrayView(int, default=ss.INT_NAN)  # This variable tracks all UIDs currently in use
-
         n = int(n_agents)
-
-        self._uid_map.grow(n)
-        self._uid_map[:] = np.arange(0, n)
-        self.uid.grow(n)
-        self.uid[:] = np.arange(0, n)
+        self.initialized = False
+        self.uid = ss.IndexArr('uid')  # This variable tracks all UIDs
+        uids = ss.uids(np.arange(n))
+        self.uid.grow(new_vals=uids)
+        self.auids = uids.copy() # This tracks all active UIDs (in practice, agents who are alive)
 
         # A slot is a special state managed internally by BasePeople
         # This is because it needs to be updated separately from any other states, as other states
         # might have fill_values that depend on the slot
-        self.slot = ss.State('slot', int, ss.INT_NAN)
-
-        self.ti = None  # Track simulation time index
-        self.dt = np.nan  # Track simulation time step
+        self.slot = ss.IndexArr('slot')
+        self.slot.grow(new_vals=uids)
 
         # User-facing collection of states
-        self.states = ss.ndict(type=ss.State)
+        self.states = ss.ndict(type=ss.Arr)
 
         # We also internally store states in a dict keyed by the memory ID of the state, so that we can have colliding names
         # e.g., across modules, but we will never change the size of a State multiple times in the same iteration over
@@ -50,12 +44,17 @@ def __init__(self, n_agents):
         # within the sim, regardless of where they are. In contrast, `People.states` offers a more user-friendly way to access
         # a selection of the states e.g., module states could be added in there, while intervention states might not
         self._states = {}
+        
         return
 
     def __len__(self):
         """ Length of people """
-        return len(self.uid)
-
+        return len(self.auids)
+    
+    @property
+    def n_uids(self):
+        return self.uid.len_used
+    
     def register_state(self, state, die=True):
         """
         Register a state with the People instance for dynamic resizing
@@ -71,60 +70,41 @@ def register_state(self, state, die=True):
             raise ValueError(errormsg)
         return
 
-    def grow(self, n, new_slots=None):
+    def grow(self, n=None, new_slots=None):
         """
         Increase the number of agents
 
         :param n: Integer number of agents to add
         :param new_slots: Optionally specify the slots to assign for the new agents. Otherwise, it will default to the new UIDs
         """
+        
+        if n is None:
+            if new_slots is None:
+                errormsg = 'Must supply either n or new_slots'
+                raise ValueError(errormsg)
+            else:
+                n = len(new_slots)
 
         if n == 0:
             return np.array([], dtype=ss.dtypes.int)
 
-        start_uid = len(self._uid_map)
-        start_idx = len(self.uid)
-
-        new_uids = np.arange(start_uid, start_uid + n)
-        new_inds = np.arange(start_idx, start_idx + n)
-
-        self._uid_map.grow(n)
-        self._uid_map[new_uids] = new_inds
-
-        self.uid.grow(n)
-        self.uid[new_inds] = new_uids
+        start_uid = self.uid.len_used
+        stop_uid = start_uid + n
+        new_uids = ss.uids(np.arange(start_uid, stop_uid))
+        self.uid.grow(new_uids, new_vals=new_uids)
 
         # We need to grow the slots as well
-        self.slot.grow(new_uids)
-        self.slot[new_uids] = new_slots if new_slots is not None else new_uids
+        new_slots = new_slots if new_slots is not None else new_uids
+        self.slot.grow(new_uids, new_vals=new_slots)
 
         for state in self._states.values():
             state.grow(new_uids)
+            
+        # Finally, update the alive indices
+        self.auids = self.auids.concat(new_uids)
 
         return new_uids
 
-    def remove(self, uids_to_remove):
-        """
-        Reduce the number of agents
-
-        :param uids_to_remove: An int/list/array containing the UID(s) to remove
-        """
-
-        # Calculate the *indices* to keep
-        keep_uids = self.uid[~np.in1d(self.uid, uids_to_remove)]  # Calculate UIDs to keep
-        keep_inds = self._uid_map[keep_uids]  # Calculate indices to keep
-
-        # Trim the UIDs and states
-        self.uid._trim(keep_inds)
-        for state in self._states.values(): # includes self.slot
-            state._trim(keep_inds)
-
-        # Update the UID map
-        self._uid_map[:] = ss.INT_NAN  # Clear out all previously used UIDs
-        self._uid_map[keep_uids] = np.arange(0, len(keep_uids))  # Assign the array indices for all of the current UIDs
-
-        return
-
     def __getitem__(self, key):
         """
         Allow people['attr'] instead of getattr(people, 'attr')
@@ -190,11 +170,11 @@ def __init__(self, n_agents, age_data=None, extra_states=None, rand_seed=0):
 
         # Handle states
         states = [
-            ss.State('age', float, np.nan), # NaN until conceived
-            ss.State('female', bool, ss.bernoulli(name='female', p=0.5)),
-            ss.State('ti_dead', int, ss.INT_NAN),  # Time index for death
-            ss.State('alive', bool, True),  # Time index for death
-            ss.State('scale', float, 1.0),
+            ss.BoolArr('alive', default=True),  # Time index for death
+            ss.BoolArr('female', default=ss.bernoulli(name='female', p=0.5)),
+            ss.FloatArr('age'), # NaN until conceived
+            ss.FloatArr('ti_dead'),  # Time index for death
+            ss.FloatArr('scale', default=1.0), # The scale factor for the agents (multiplied for making results)
         ]
         states.extend(sc.promotetolist(extra_states))
 
@@ -222,7 +202,6 @@ def get_age_dist(age_data):
             age_props = age_props / age_props.sum()
             return ss.choice(a=age_bins, p=age_props)
 
-
     def initialize(self, sim):
         """ Initialization """
 
@@ -233,8 +212,8 @@ def initialize(self, sim):
         
         # For People initialization, first initialize slots, then initialize RNGs, then initialize remaining states
         # This is because some states may depend on RNGs being initialized to generate initial values
-        self.slot.initialize(sim)
-        self.slot[:] = self.uid
+        self.uid.set_people(sim.people)
+        self.slot.set_people(sim.people)
 
         # Initialize states
         # Age is handled separately because the default value for new agents is NaN until they are concieved/born whereas
@@ -297,13 +276,15 @@ def remove_dead(self, sim):
         """
         Remove dead agents
         """
-        uids_to_remove = ss.true(self.dead)
-        if len(uids_to_remove):
-            self.remove(uids_to_remove)
-
-        # Remove the UIDs from the network too
-        for network in sim.networks.values():
-            network.remove_uids(uids_to_remove)
+        uids = (~self.alive).uids
+        if len(uids):
+            
+            # Remove the UIDs from the networks too
+            for network in sim.networks.values():
+                network.remove_uids(uids) # TODO: only run once every nth timestep
+                
+            # Calculate the indices to keep
+            self.auids = self.auids.remove(uids)
 
         return
 
@@ -314,7 +295,7 @@ def update_post(self, sim):
         :param sim:
         :return:
         """
-        self.age[self.alive] += self.dt
+        self.age[self.alive.uids] += self.dt
         return
 
     def resolve_deaths(self):
@@ -323,7 +304,7 @@ def resolve_deaths(self):
 
         :return:
         """
-        death_uids = ss.true(self.ti_dead <= self.ti)
+        death_uids = (self.ti_dead <= self.ti).uids
         self.alive[death_uids] = False
         return death_uids
 
@@ -338,7 +319,7 @@ def male(self):
         return ~self.female
 
     @property
-    def f(self):
+    def f(self): # TODO: remove?
         """ Shorthand for female """
         return self.female
 
@@ -349,17 +330,18 @@ def m(self):
 
     def init_results(self, sim):
         sim.results += [
-            ss.Result(None, 'n_alive', sim.npts, ss.dtypes.int, scale=True),
+            ss.Result(None, 'n_alive',    sim.npts, ss.dtypes.int, scale=True),
             ss.Result(None, 'new_deaths', sim.npts, ss.dtypes.int, scale=True),
             ss.Result(None, 'cum_deaths', sim.npts, ss.dtypes.int, scale=True),
         ]
         return
 
     def update_results(self, sim):
+        ti = sim.ti
         res = sim.results
-        res.n_alive[self.ti] = np.count_nonzero(self.alive)
-        res.new_deaths[self.ti] = np.count_nonzero(self.ti_dead == self.ti)
-        res.cum_deaths[self.ti] = np.sum(res.new_deaths[:sim.ti])
+        res.n_alive[ti] = np.count_nonzero(self.alive)
+        res.new_deaths[ti] = np.count_nonzero(self.ti_dead == ti)
+        res.cum_deaths[ti] = np.sum(res.new_deaths[:ti]) # TODO: inefficient to compute the cumulative sum on every timestep!
         return
 
     def request_death(self, uids):
@@ -387,10 +369,5 @@ def request_death(self, uids):
         :param uids: Agent IDs to request deaths for
         :return: UIDs of agents that have been scheduled to die on this timestep
         """
-
-        # Only update the time of death for agents that are currently alive. This way modules cannot
-        # modify the time of death for agents that have already died. Noting that if remove_people is
-        # enabled then often such agents would not be present in the simulation anyway
-        uids = ss.true(self.alive[uids])
         self.ti_dead[uids] = self.ti
         return
diff --git a/starsim/products.py b/starsim/products.py
index b5afd79c..a49987b0 100644
--- a/starsim/products.py
+++ b/starsim/products.py
@@ -59,7 +59,8 @@ def administer(self, sim, uids, return_format='dict'):
         for disease in self.diseases:
             for state in self.health_states:
                 this_state = getattr(sim.diseases[disease], state)
-                these_uids = ss.true(this_state[uids])
+                true_uids = this_state.uids # Find people for which this state is true
+                these_uids = true_uids.intersect(uids) # Find intersection of people in this state and the supplied UIDs
 
                 # Filter the dataframe to extract test results for people in this state
                 df_filter = (self.df.state == state) & (self.df.disease == disease)
@@ -68,8 +69,8 @@ def administer(self, sim, uids, return_format='dict'):
                 self.result_dist.pk = probs  # Overwrite distribution probabilities
 
                 # Sort people into one of the possible result states and then update their overall results
-                this_result = self.result_dist.rvs(these_uids)-these_uids # TODO: check!
-                row_inds = results.uids.isin(these_uids)
+                this_result = self.result_dist.rvs(these_uids)
+                row_inds = sc.findinds(results.uids.isin(these_uids))
                 results.loc[row_inds, 'result'] = np.minimum(this_result, results.loc[row_inds, 'result'])
 
             if return_format == 'dict':
@@ -106,7 +107,8 @@ def administer(self, sim, uids, return_format='dict'):
             for state in self.health_states:
 
                 pre_tx_state = getattr(disease, state)
-                these_uids = ss.true(pre_tx_state[uids])
+                true_uids = pre_tx_state.uids # People in this state
+                these_uids = true_uids.intersect(uids)
 
                 if len(these_uids):
 
@@ -141,8 +143,8 @@ def administer(self, sim, uids, return_format='dict'):
 class Vx(Product):
     """ Vaccine product """
     def __init__(self, diseases=None, pars=None, par_dists=None, *args, **kwargs):
-        pars = ss.omerge({}, pars)
-        par_dists = ss.omerge({}, par_dists)
+        pars = ss.dictmerge({}, pars)
+        par_dists = ss.dictmerge({}, par_dists)
         super().__init__(pars, par_dists, *args, **kwargs)
         self.diseases = sc.tolist(diseases)
         return
diff --git a/starsim/settings.py b/starsim/settings.py
index d0bc8103..1bccd52e 100644
--- a/starsim/settings.py
+++ b/starsim/settings.py
@@ -9,9 +9,7 @@
 import numpy as np
 import sciris as sc
 
-__all__ = ['INT_NAN', 'dtypes', 'options']
-
-INT_NAN = 2147483647 # From np.iinfo(np.int32).max: value to use to flag invalid content (i.e., an integer value we are treating like NaN, since NaN can't be stored in an integer array)
+__all__ = ['dtypes', 'options']
 
 # Define Starsim-default data types
 dtypes = sc.objdict(
diff --git a/starsim/sim.py b/starsim/sim.py
index 6d3c89f9..de7e55c4 100644
--- a/starsim/sim.py
+++ b/starsim/sim.py
@@ -467,10 +467,6 @@ def step(self):
         # Advance random number generators forward to prepare for any random number calls that may be necessary on this step
         self.dists.jump(to=self.ti+1)  # +1 offset because ti=0 is used on initialization
 
-        # Clean up dead agents, if removing agents is enabled
-        if self.pars.remove_dead and (self.ti % self.pars.remove_dead == 0):
-            self.people.remove_dead(self)
-
         # Update demographic modules (create new agents from births/immigration, schedule non-disease deaths and emigration)
         for dem_mod in self.demographics.values():
             dem_mod.update(self)
@@ -511,6 +507,9 @@ def step(self):
 
         for analyzer in self.analyzers.values():
             analyzer(self)
+            
+        # Clean up dead agents
+        self.people.remove_dead(self)
 
         # Tidy up
         self.ti += 1
diff --git a/starsim/states.py b/starsim/states.py
index 44208ebd..4dfa48ae 100644
--- a/starsim/states.py
+++ b/starsim/states.py
@@ -3,14 +3,15 @@
 """
 
 import numpy as np
-import sciris as sc
-import numba as nb
 import starsim as ss
-from starsim.settings import INT_NAN
-from starsim.settings import dtypes as sdt
-from numpy.lib.mixins import NDArrayOperatorsMixin  # Inherit from this to automatically gain operators like +, -, ==, <, etc.
 
-__all__ = ['check_dtype', 'UIDArray', 'State', 'ArrayView']
+# Shorten these for performance
+ss_float = ss.dtypes.float
+ss_int   = ss.dtypes.int
+ss_bool  = ss.dtypes.bool
+
+__all__ = ['check_dtype', 'Arr', 'FloatArr', 'BoolArr', 'IndexArr', 'uids']
+
 
 def check_dtype(dtype, default=None):
     """ Check that the supplied dtype is one of the supported options """
@@ -24,11 +25,11 @@ def check_dtype(dtype, default=None):
             dtype = type(default)
     
     if dtype in ['float', float, np.float64, np.float32]:
-        dtype = sdt.float
+        dtype = ss_float
     elif dtype in ['int', int, np.int64, np.int32]:
-        dtype = sdt.int
+        dtype = ss_int
     elif dtype in ['bool', bool, np.bool_]:
-        dtype = sdt.bool
+        dtype = ss_bool
     else:
         warnmsg = f'Data type {type(default)} not a supported data type; set warn=False to suppress warning'
         ss.warn(warnmsg)
@@ -36,450 +37,207 @@ def check_dtype(dtype, default=None):
     return dtype
 
 
-class UIDArray(NDArrayOperatorsMixin):
+class Arr(np.lib.mixins.NDArrayOperatorsMixin):
     """
-    This is a class that allows indexing by UID but does not support dynamic growth
-    It's kind of like a Pandas series but one that only supports a monotonically increasing
-    unique integer index, and that we can customize and optimize indexing for.
+    Store a state of the agents (e.g. age, infection status, etc.) as an array.
+    
+    In practice, ``Arr`` objects can be used interchangeably with NumPy arrays.
+    They have two main data interfaces: ``Arr.raw`` contains the "raw", underlying
+    NumPy array of the data. ``Arr.values`` contains the "active" values, which
+    usually corresponds to agents who are alive.
+    
+    By default, operations are performed on active agents only (specified by ``Arr.auids``,
+    which is a pointer to ``sim.people.auids``). For example, ``sim.people.age.mean()``
+    will only use the ages of active agents. Thus, ``sim.people.age.mean()``
+    is equal to ``sim.people.age.values.mean()``, not ``sim.people.age.raw.mean()``.
     
-    We explictly do NOT support slicing (except for `[:]`), as these arrays are indexed by UID and slicing
-    by UID can be confusing/ambiguous when there are missing values. Indexing with a list/array returns
-    another UIDArray instance which enables chained filtering
+    If indexing by an int or slice, ``Arr.values`` is used. If indexing by an
+    ``ss.uids`` object, ``Arr.raw`` is used. ``Arr`` objects can't be directly
+    indexed by a list or array of ints, as this would be ambiguous about whether
+    ``values`` or ``raw`` is intended. For example, if there are 1000 people in a 
+    simulation and 100 of them have died, ``sim.people.age[999]`` will return
+    an ``IndexError`` (since ``sim.people.age[899]`` is the last active agent),
+    whereas ``sim.people.age[ss.uids(999)]`` is valid.
+
+    Args: 
+        name (str): The name for the state (also used as the dictionary key, so should not have spaces etc.)
+        dtype (class): The dtype to use for this instance (if None, infer from value)
+        default (any): Specify default value for new agents. This can be
+        - A scalar with the same dtype (or castable to the same dtype) as the State
+        - A callable, with a single argument for the number of values to produce
+        - A ``ss.Dist`` instance
+        nan (any): the value to use to represent NaN (not a number); also used as the default value if not supplied
+        raw (arr): if supplied, the raw values to use
+        label (str): The human-readable name for the state
+        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)
     """
-
-    __slots__ = ('values', '_uid_map', 'uid')
-
-    def __init__(self, values=None, uid=None, uid_map=None):
-
-        self.values = values
-        self.uid = uid
-
-        if uid_map is None and uid is not None:
-            # Construct a local UID map as opposed to using a shared one (i.e., the one for all agents contained in the People instance)
-            self._uid_map = np.full(np.max(uid) + 1, fill_value=ss.INT_NAN, dtype=sdt.int)
-            self._uid_map[uid] = np.arange(len(uid))
-        else:
-            self._uid_map = uid_map
-        return
-
-    def __repr__(self):
-        if len(self) == 0:
-            return f'<{self.__class__.__name__} (empty)>'
-        df = self.to_df()
-        return df.__repr__()
-    
-    def to_df(self):
-        """ Convert to a dataframe """
-        df = sc.dataframe({'Quantity':self.values.T}, index=self.uid)
-        df.index.name = 'UID'
-        return df
-
-    @property
-    def dtype(self):
-        return self.values.dtype
-
-    @staticmethod
-    @nb.njit(cache=True)
-    def _get_vals_uids(vals, key, uid_map):
-        """
-        Extract values from a collection of UIDs
-
-        This function is used to retrieve values based on UID. As indexing a UIDArray returns a new UIDArray,
-        this method also populates the new UID map for use in the subsequently created UIDArray, avoiding the
-        need to re-compute it separately.
-
-        Args:
-            vals: A 1D np.ndarray containing the values
-            key: A 1D np.ndnarray of integers containing the UIDs to query
-            uid_map: A 1D np.ndarray of integers mapping UID to array position in ``vals``
+    def __init__(self, name, dtype=None, default=None, nan=None, raw=None, label=None, coerce=True, skip_init=False):
+        if coerce:
+            dtype = check_dtype(dtype, default)
         
-        Returns:
-            A tuple of (values, uids, new_uid_map) suitable for passing into the UIDArray constructor
-        """
-        out = np.empty(len(key), dtype=vals.dtype)
-        new_uid_map = np.full(uid_map.shape[0], fill_value=INT_NAN, dtype=np.int64)
-
-        for i,kv in enumerate(key):
-            idx = uid_map[kv]
-            if idx == INT_NAN:
-                raise IndexError('UID not present in array')
-            out[i] = vals[idx]
-            new_uid_map[kv] = i
-        return out, key, new_uid_map
-    
-    @staticmethod
-    @nb.njit(cache=True, parallel=True)
-    def _get_vals_uids_par(vals, key, uid_map):
-        """
-        Parallelized version of _get_vals_uids() -- need both because this is much slower for small arrays
-        """
-        len_key = len(key)
-        out = np.empty(len_key, dtype=vals.dtype)
-        new_uid_map = np.full(uid_map.shape[0], fill_value=INT_NAN, dtype=np.int64)
-
-        for i in nb.prange(len_key):
-            kv = key[i]
-            idx = uid_map[kv]
-            out[i] = vals[idx]
-            new_uid_map[kv] = i
-        return out, key, new_uid_map
-
-    @staticmethod
-    @nb.njit(cache=True)
-    def _set_vals_uids_multiple(vals, key, uid_map, value):
-        """
-        Insert an array of values based on UID
-
-        Args:
-            vals: A reference to a 1D np.ndarray in which to insert the values
-            key: A 1D np.ndnarray of integers containing the UIDs to add values for
-            uid_map:  A 1D np.ndarray of integers mapping UID to array position in ``vals``
-            value: A 1D np.ndarray the same length as ``key`` containing values to insert
-        """
-
-        for i,kv in enumerate(key):
-            if kv >= len(uid_map):
-                errormsg = f'UID not present in array (requested UID ({key[i]}) is larger than the maximum UID in use ({len(uid_map)}))'
-                raise IndexError(errormsg)
-            idx = uid_map[kv]
-            if idx == INT_NAN:
-                raise IndexError('UID not present in array')
-            elif idx >= len(vals):
-                errormsg = f'Attempted to write to a non-existant index ({idx}) - this can happen if attempting to write to new entries that have not yet been allocated via grow()'
-                raise IndexError(errormsg)
-            vals[idx] = value[i]
+        # Set attributes
+        self.name = name
+        self.label = label or name
+        self.default = default
+        self.nan = nan
+        self.dtype = dtype
+        
+        # Properties that are initialized later
+        self.raw = np.empty(0, dtype=dtype)
+        self.people = None
+        self.len_used = 0
+        self.len_tot = 0
+        self.initialized = skip_init
+        if raw is not None:
+            self.grow(new_uids=uids(np.arange(len(raw))), new_vals=raw)
         return
-
-    @staticmethod
-    @nb.njit(cache=True)
-    def _set_vals_uids_single(vals, key, uid_map, value):
+    
+    def __repr__(self):
+        arr_str = np.array2string(self.values, max_line_width=200)
+        string = f'<{self.__class__.__name__} "{str(self.name)}", len={len(self)}, {arr_str}>'
+        return string
+    
+    def __len__(self):
+        return len(self.auids)
+    
+    def _convert_key(self, key):
         """
-        Insert a single value into multiple UIDs
-
-        Args:
-            vals: A reference to a 1D np.ndarray in which to insert the values
-            key: A 1D np.ndnarray of integers containing the UIDs to add values for
-            uid_map:  A 1D np.ndarray of integers mapping UID to array position in ``vals``
-            value: A scalar value to insert at every position specified by ``key``
+        Used for getitem and setitem to determine whether the key is indexing
+        the raw array (``raw``) or the active agents (``values``), and to convert
+        the key to array indices if needed.
         """
-        for i,kv in enumerate(key):
-            if kv >= len(uid_map):
-                raise IndexError('UID not present in array (requested UID is larger than the maximum UID in use)')
-            idx = uid_map[kv]
-            if idx == INT_NAN:
-                raise IndexError('UID not present in array')
-            elif idx >= len(vals):
-                raise Exception('Attempted to write to a non-existant index - this can happen if attempting to write to new entries that have not yet been allocated via grow()')
-            vals[idx] = value
-
+        use_raw = True
+        if isinstance(key, uids):
+            pass
+        elif isinstance(key, (BoolArr, IndexArr)):
+            key = key.uids
+        elif isinstance(key, (slice, int)):
+            use_raw = False
+        elif not np.isscalar(key) and len(key) == 0: # Handle [], np.array([]), etc.
+            key = uids()
+        else:
+            errormsg = f'Indexing an Arr ({self.name}) by ({key}) is ambiguous or not supported. Use ss.uids() instead, or index Arr.raw or Arr.values.'
+            raise Exception(errormsg)
+        
+        return key, use_raw
+    
     def __getitem__(self, key):
-        if np.iterable(key) and len(key) > 10_000: # Approximate cutoff for when the overhead becomes worthwhile
-            gvu_func = self._get_vals_uids_par
+        key, use_raw = self._convert_key(key)
+        if use_raw:
+            return self.raw[key]
         else:
-            gvu_func = self._get_vals_uids
-            
-        try:
-            if isinstance(key, (int, np.integer)):
-                # Handle getting a single item by UID
-                return self.values[self._uid_map[key]]
-            elif isinstance(key, (np.ndarray, UIDArray, ArrayView)):
-                if key.dtype.kind == 'b':
-                    # Handle accessing items with a logical array. Surprisingly, it seems faster to use nonzero() to convert
-                    # it to indices first. Also, the pure Python implementation is difficult to improve upon using numba
-                    mapped_key = key.__array__().nonzero()[0]
-                    uids = self.uid.__array__()[mapped_key]
-                    new_uid_map = np.full(len(self._uid_map), fill_value=INT_NAN, dtype=int)
-                    new_uid_map[uids] = np.arange(len(uids))
-                    values = self.values[mapped_key]
-                else:
-                    # Access items by an array of integers. We do get a decent performance boost from using numba here
-                    values, uids, new_uid_map = gvu_func(self.values, key.__array__(), self._uid_map.__array__())
-            elif isinstance(key, slice):
-                if key.start is None and key.stop is None and key.step is None:
-                    return sc.dcp(self)
-                else:
-                    raise Exception('Slicing not supported - slice the .values attribute by index instead e.g., x.values[0:5], not x[0:5]')
-            else:
-                # This branch is specifically handling the user passing in a list of integers instead of an array, therefore
-                # we need an additional conversion to an array first using np.fromiter to improve numba performance
-                values, uids, new_uid_map = gvu_func(self.values, np.fromiter(key, dtype=int), self._uid_map.__array__())
-            return UIDArray(values=values, uid=uids, uid_map=new_uid_map)
-        except IndexError as e:
-            if str(INT_NAN) in str(e):
-                raise IndexError(f'UID {key} not present in array')
-            else:
-                raise e
-
+            return self.values[key]
+    
     def __setitem__(self, key, value):
-        """
-        NB: the use of .__array__() calls is to access the array interface and thereby treat both np.ndarray and ArrayView instances
-        in the same way without needing an additional type check. This is also why the UIDArray.dtype property is defined. Noting
-        that for a State, the uid_map is a dynamic view attached to the People, but after an indexing operation, it will be a bare
-        UIDArray that has an ordinary numpy array as the uid_map.
-        """
-        try:
-            if isinstance(key, (int, np.integer)):
-                return self.values.__setitem__(self._uid_map[key], value)
-            elif isinstance(key, (np.ndarray, UIDArray)):
-                if key.dtype.kind == 'b':
-                    self.values.__setitem__(key.__array__().nonzero()[0], value)
-                else:
-                    if isinstance(value, (np.ndarray, UIDArray)):
-                        return self._set_vals_uids_multiple(self.values, key, self._uid_map.__array__(), value.__array__())
-                    else:
-                        return self._set_vals_uids_single(self.values, key, self._uid_map.__array__(), value)
-            elif isinstance(key, slice):
-                if key.start is None and key.stop is None and key.step is None:
-                    return self.values.__setitem__(key, value)
-                else:
-                    raise Exception('Slicing not supported - slice the .values attribute by index instead e.g., x.values[0:5], not x[0:5]')
-            else:
-                if isinstance(value, (np.ndarray, UIDArray)):
-                    return self._set_vals_uids_multiple(self.values, np.fromiter(key, dtype=int), self._uid_map.__array__(), value.__array__())
-                else:
-                    return self._set_vals_uids_single(self.values, np.fromiter(key, dtype=int), self._uid_map.__array__(), value)
-        except IndexError as e:
-            if str(INT_NAN) in str(e):
-                raise IndexError(f'UID {key} not present in array')
-            else:
-                raise e
-
+        key, use_raw = self._convert_key(key)
+        if use_raw:
+            self.raw[key] = value
+        else:
+            self.raw[self.auids[key]] = value
+            
     def __getattr__(self, attr):
         """ Make it behave like a regular array mostly -- enables things like sum(), mean(), etc. """
         if attr in ['__deepcopy__', '__getstate__', '__setstate__']:
             return self.__getattribute__(attr)
         else:
             return getattr(self.values, attr)
-
-    # Make it behave like a regular array mostly
-    def __len__(self):
-        return len(self.values)
-
-    def __contains__(self, *args, **kwargs):
-        return self.values.__contains__(*args, **kwargs)
-
-    def astype(self, *args, **kwargs):
-        return UIDArray(values=self.values.astype(*args, **kwargs), uid=self.uid, uid_map=self._uid_map)
-
-    def sum(self, *args, **kwargs):
-        return self.values.sum(*args, **kwargs)
-
-    def mean(self, *args, **kwargs):
-        return self.values.mean(*args, **kwargs)
-
-    def any(self, *args, **kwargs):
-        return self.values.any(*args, **kwargs)
-
-    def all(self, *args, **kwargs):
-        return self.values.all(*args, **kwargs)
-
-    def count(self, *args, **kwargs):
-        return np.count_nonzero(self.values, *args, **kwargs)
-
+        
+    def __gt__(self, other): return self.asnew(self.values > other,  cls=BoolArr)
+    def __lt__(self, other): return self.asnew(self.values < other,  cls=BoolArr)
+    def __ge__(self, other): return self.asnew(self.values >= other, cls=BoolArr)
+    def __le__(self, other): return self.asnew(self.values <= other, cls=BoolArr)
+    def __eq__(self, other): return self.asnew(self.values == other, cls=BoolArr)
+    def __ne__(self, other): return self.asnew(self.values != other, cls=BoolArr)
+    
+    def __and__(self, other): raise BooleanOperationError(self)
+    def __or__(self, other):  raise BooleanOperationError(self)
+    def __xor__(self, other): raise BooleanOperationError(self)
+    def __invert__(self):     raise BooleanOperationError(self)
+    
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        """ For almost everything else, behave like a normal NumPy array on Arr.values """
+        inputs = [x.values if isinstance(x, Arr) else x for x in inputs]
+        return getattr(ufunc, method)(*inputs, **kwargs)
+    
     @property
-    def shape(self):
-        return self.values.shape
+    def auids(self):
+        """ Link to the indices of active agents -- sim.people.auids """
+        try:
+            return self.people.auids
+        except:
+            print('Non-initialized States object')
+            return uids(np.arange(len(self.raw)))
+    
+    def count(self):
+        return np.count_nonzero(self.values)
 
     @property
-    def __array_interface__(self):
-        return self.values.__array_interface__
-
-    def __array__(self):
-        return self.values
-
-    def __array_ufunc__(self, *args, **kwargs):
-        if args[1] != '__call__':
-            # This is a generic catch-all for ufuncs that are not being applied with '__call__' (e.g., operations returning a scalar like 'np.sum()' use reduce instead)
-            args = [(x if x is not self else self.values) for x in args]
-            kwargs = {k: v if v is not self else self.values for k, v in kwargs.items()}
-            return self.values.__array_ufunc__(*args, **kwargs)
-
-        args = [(x if x is not self else self.values) for x in args]
-        if 'out' in kwargs and kwargs['out'][0] is self:
-            # In-place operations like += use this branch
-            kwargs['out'] = self.values
-            args[0](*args[2:], **kwargs)
-            return self
-        else:
-            out = args[0](*args[2:], **kwargs)
-            if isinstance(out, UIDArray):
-                # For some operations (e.g., those involving two UIDArrays) the result of the ufunc will already be a UIDArray
-                # In particular, operating on two states will result in a UIDArray where the references to the original People uid_map and uids
-                # are still intact. In such cases, we can return the resulting UIDArray directly
-                return out
+    def values(self):
+        """ Return the values of the active agents """
+        return self.raw[self.auids]
+
+    def set(self, uids, new_vals=None):
+        """ Set the values for the specified UIDs"""
+        if new_vals is None: 
+            if isinstance(self.default, ss.Dist):
+                new_vals = self.default.rvs(uids)
+            elif callable(self.default):
+                new_vals = self.default(len(uids))
+            elif self.default is not None:
+                new_vals = self.default
             else:
-                # Otherwise, if the result of the ufunc is an array (e.g., because one of the arguments was an array) then
-                # we need to wrap it up in a new UIDArray. With '__call__' the dimensions should hopefully be the same and we should
-                # be able to reuse the UID arrays directly
-                return UIDArray(values=out, uid=self.uid, uid_map=self._uid_map)
-
-    def __array_wrap__(self, out_arr, context=None):
-        # This allows numpy operations addition etc. to return instances of UIDArray
-        if out_arr.ndim == 0:
-            return out_arr.item()
-        return UIDArray(values=out_arr, uid=self.uid, uid_map=self._uid_map) # Hardcoding class means State can inherit from UIDArray but return UIDArray base instances
-
-
-class ArrayView(NDArrayOperatorsMixin):
-
-    __slots__ = ('_data', '_view', 'n', 'default')
-
-    def __init__(self, dtype, default=None, coerce=True):
-        """
-        Args:
-            dtype (class): The dtype to use for this instance (if None, infer from value)
-            default (any): Specify default value for new agents. This can be
-            coerce (bool): Whether to ensure the the data is one of the supported data types
-        """
-        if coerce:
-            dtype = check_dtype(dtype, default)
-        self.default = default if default is not None else dtype()
-        self.n = 0  # Number of agents currently in use
-        self._data = np.empty(0, dtype=dtype)  # The underlying memory array (length at least equal to n)
-        self._view = None  # The view corresponding to what is actually accessible (length equal to n)
-        self._map_arrays()
-        return
-
-    @property
-    def _s(self):
-        """ Return the size of the underlying array (maximum number of agents that can be stored without reallocation) """
-        return len(self._data)
-
-    @property
-    def dtype(self):
-        """
-        The specified dtype and the underlying array dtype can be different. For instance, the user might pass in
-        ArrayView(dtype=int) but the underlying array's dtype will be np.dtype('int32'). This distinction is important
-        because the numpy dtype has attributes like 'kind' that the input dtype may not have. We need the ArrayView's
-        dtype to match that of the underlying array so that it can be more seamlessly exchanged with direct numpy arrays
-        Therefore, we retain the original dtype in ArrayView._dtype() and use
-        """
-        return self._data.dtype
-
-    def __len__(self):
-        """ Return the number of active elements """
-        return self.n
-
-    def __repr__(self):
-        """ Print out the numpy view directly """
-        return self._view.__repr__()
-
-    def grow(self, n):
-        """ If the total number of agents exceeds the array size, extend the underlying arrays """
-        if self.n + n > self._s:
-            n_new = max(n, int(self._s / 2))  # Minimum 50% growth
-            self._data = np.concatenate([self._data, np.full(n_new, dtype=self.dtype, fill_value=self.default)], axis=0)
-        self.n += n  # Increase the count of the number of agents by `n` (the requested number of new agents)
-        self._map_arrays()
-        return
-
-    def _trim(self, inds):
-        """ Keep only specified indices """
-        # Note that these are indices, not UIDs!
-        n = len(inds)
-        self._data[:n] = self._data[inds]
-        self._data[n:self.n] = self.default
-        self.n = n
-        self._map_arrays()
-        return
-
-    def __getstate__(self):
-        """ When pickling, skip storing the `._view` attribute, which should be re-linked after unpickling """
-        return {k:getattr(self, k) for k in self.__slots__ if k != '_view'}
-
-    def __setstate__(self, state):
-        """ Re-map arrays after unpickling so that `.view` is a reference to the correct array in-memory """
-        for k,v in state.items():
-            setattr(self, k, v)
-        self._map_arrays()
+                new_vals = self.nan
+        self.raw[uids] = new_vals
+        return new_vals
+    
+    def set_nan(self, uids):
+        """ Shortcut function to set values to NaN """
+        self.raw[uids] = self.nan
         return
-
-    def _map_arrays(self):
-        """
-        Set main simulation attributes to be views of the underlying data
-
-        This method should be called whenever the number of agents required changes
-        (regardless of whether or not the underlying arrays have been resized)
+    
+    def grow(self, new_uids=None, new_vals=None):
         """
-        self._view = self._data[:self.n]
-        return
-
-    def __getitem__(self, key):
-        return self._view.__getitem__(key)
-
-    def __setitem__(self, key, value):
-        self._view.__setitem__(key, value)
-        return
-        
-    def __getattr__(self, attr):
-        """ Make it behave like a regular array mostly -- enables things like sum(), mean(), etc. """
-        if attr in ['__deepcopy__', '__getstate__', '__setstate__']:
-            return self.__getattribute__(attr)
-        else:
-            return getattr(self._view, attr)
-
-    @property
-    def __array_interface__(self):
-        return self._view.__array_interface__
-
-    def __array__(self):
-        return self._view
+        Add new agents to an Arr
 
-    def __array_ufunc__(self, *args, **kwargs):
-        args = [(x if x is not self else self._view) for x in args]
-        kwargs = {k: v if v is not self else self._view for k, v in kwargs.items()}
-        return self._view.__array_ufunc__(*args, **kwargs)
-
-
-class State(UIDArray):
-
-    __slots__ = ('values', '_uid_map', 'uid', 'default', 'name', 'label', '_data', 'values', '_initialized')
-
-    def __init__(self, name, dtype=None, default=None, label=None, coerce=True):
-        """
-        Store a state of the agents (e.g. age, infection status, etc.)
+        This method is normally only called via `People.grow()`.
 
-        Args: 
-            name (str): The name for the state (also used as the dictionary key, so should not have spaces etc.)
-            dtype (class): The dtype to use for this instance (if None, infer from value)
-            default (any): Specify default value for new agents. This can be
-            - A scalar with the same dtype (or castable to the same dtype) as the State
-            - A callable, with a single argument for the number of values to produce
-            - A ``ss.Dist`` instance
-            label (str): The human-readable name for the state
-            coerce (bool): Whether to ensure the the data is one of the supported data types
+        Args:
+            new_uids: Numpy array of UIDs for the new agents being added
+            new_vals: If provided, assign these state values to the new UIDs
         """
-        super().__init__()  # Call the UIDArray constructor
-        
-        if coerce:
-            dtype = check_dtype(dtype, default)
+        if new_uids is None and new_vals is not None: # Used as a shortcut to avoid needing to supply twice
+            new_uids = new_vals
+        orig_len = self.len_used
+        n_new = len(new_uids)
+        self.len_used += n_new  # Increase the count of the number of agents by `n` (the requested number of new agents)
         
-        if default is None:
-            default = dtype()
+        # Physically reshape the arrays, if needed
+        if orig_len + n_new > self.len_tot:
+            n_grow = max(n_new, self.len_tot//2)  # Minimum 50% growth, since growing arrays is slow
+            new_empty = np.empty(n_grow, dtype=self.dtype) # 10x faster than np.zeros()
+            self.raw = np.concatenate([self.raw, new_empty], axis=0)
+            self.len_tot = len(self.raw)
+            if n_grow > n_new: # We added extra space at the end, set to NaN
+                nan_uids = np.arange(self.len_used, self.len_tot)
+                self.set_nan(nan_uids)
         
-        # Set attributes
-        self.default = default
-        self.name = name
-        self.label = label or name
-        self._data = ArrayView(dtype=dtype)
-        self.values = self._data._view
-        self._initialized = False
+        # Set new values, and NaN if needed
+        self.set(new_uids, new_vals=new_vals) # Assign new default values to those agents
         return
-
-    def __repr__(self):
-        if not self._initialized:
-            return f'<State {self.name} (uninitialized)>'
-        else:
-            return UIDArray.__repr__(self)
-
-    def _new_vals(self, uids):
-        if isinstance(self.default, ss.Dist):
-            new_vals = self.default.rvs(uids)
-        elif callable(self.default):
-            new_vals = self.default(len(uids))
+    
+    def set_people(self, people):
+        """ Reset the people object associated with this state """
+        if isinstance(people, ss.People): # It's people, it's fine
+            pass
+        elif isinstance(people, ss.Sim): # Actually a sim
+            people = people.people
         else:
-            new_vals = self.default
-        return new_vals
+            errormsg = f'Must supply a Sim or People object, not {type(people)}'
+            raise TypeError(errormsg)
+        assert people.initialized, 'People must be initialized before initializing states'
+        self.people = people # Shorten since used a lot
+        return
 
     def initialize(self, sim):
         """
@@ -501,59 +259,139 @@ def initialize(self, sim):
         Args:
             sim: A `Sim` instance that contains an initialized `People` object
         """
-
-        if self._initialized:
+        # Skip if already initialized
+        if self.initialized:
             return
 
+        # Establish connection with the People object
         people = sim.people
-        assert people.initialized, 'People must be initialized before initializing states'
+        self.set_people(people)
+        people.register_state(self)
         
         # Connect any distributions in the default to RNGs in the Sim
         if isinstance(self.default, ss.Dist):
             self.default.initialize(module=self, sim=sim)
 
-        # Establish connection with the People object
-        people.register_state(self)
-        self._uid_map = people._uid_map
-        self.uid = people.uid
-
         # Populate initial values
-        self._data.grow(len(self.uid))
-        self._data[:len(self.uid)] = self._new_vals(self.uid)
-        self.values = self._data._view
-
-        self._initialized = True
+        self.grow(people.uid)
+        self.initialized = True
         return
 
-    def grow(self, uids):
-        """
-        Add state for new agents
-
-        This method is normally only called via `People.grow()`.
+    def asnew(self, arr=None, cls=None):
+        """ Duplicate and copy (rather than link) data, optionally resetting the array """
+        if cls is None:
+            cls = self.__class__
+        if arr is None:
+            arr = self.values
+        new = object.__new__(cls) # Create a new Arr instance
+        new.__dict__ = self.__dict__.copy() # Copy pointers
+        new.dtype = arr.dtype # Set to correct dtype
+        new.raw = np.empty_like(new.raw, dtype=new.dtype) # Copy values, breaking reference
+        new.raw[new.auids] = arr
+        return new
+
+
+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)
+        return
+    
+    @property
+    def isnan(self):
+        """ Return indices that are NaN """
+        return np.nonzero(np.isnan(self.values))[0]
 
-        Args:
-            uids: Numpy array of UIDs for the new agents being added This array should have length n
-        """
+    @property
+    def notnan(self):
+        """ Return indices that are not-NaN """
+        return np.nonzero(~np.isnan(self.values))[0]
+    
+    @property
+    def notnanvals(self):
+        """ Return values that are not-NaN """
+        vals = self.values # Shorten and avoid double indexing
+        out = vals[np.nonzero(~np.isnan(vals))[0]]
+        return out
 
-        n = len(uids)
-        self._data.grow(n)
-        self.values = self._data._view
-        self._data[-n:] = self._new_vals(uids)
+    
+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)
         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)
+    
+    @property
+    def uids(self):
+        """ Convert True values to UIDs """
+        return self.auids[np.nonzero(self.values)[0]]
 
-    def _trim(self, inds):
-        # Trim arrays to remove agents - should only be called via `People.remove()`
-        self._data._trim(inds)
-        self.values = self._data._view
+    
+class IndexArr(Arr):
+    """ A special class of IndexArr used for UIDs and RNG IDs """
+    def __init__(self, name, label=None):
+        super().__init__(name=name, dtype=ss_int, default=None, nan=-1, label=label, coerce=False, skip_init=True)
+        self.raw = uids(self.raw)
         return
+    
+    @property
+    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]
 
-    def __getstate__(self):
-        # When pickling, skip storing the `.values` attribute, which should be re-linked after unpickling
-        return {k:getattr(self, k) for k in self.__slots__ if k != 'values'}
-
-    def __setstate__(self, state):
-        # When unpickling, re-link the `.values` attribute
-        for k,v in state.items():
-            setattr(self, k, v)
-        self.values = self._data._view
+    @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 """
+        super().grow(new_uids=new_uids, new_vals=new_vals)
+        self.raw = uids(self.raw)
         return
+    
+    
+class uids(np.ndarray):
+    """
+    Class to specify that integers should be interpreted as UIDs.
+    
+    For all practical purposes, behaves like a NumPy integer array. However,
+    has additional methods ``uids.concat()`` (instance method), ``ss.uids.cat()``
+    (class method), ``uids.remove()``, and ``uids.intersect()`` to simplify common
+    UID operations.    
+    """
+    def __new__(cls, arr=None):
+        if arr is None:
+            arr = np.empty(0, dtype=ss_int)
+        return np.asarray(arr).view(cls)
+    
+    def concat(self, other, **kw): # TODO: why can't they both be called cat()?
+        """ Equivalent to np.concatenate(), but return correct type """
+        return np.concatenate([self, other], **kw).view(self.__class__)
+    
+    @classmethod
+    def cat(cls, *args, **kw):
+        """ Equivalent to np.concatenate(), but return correct type """
+        arrs = args[0] if len(args) == 1 else args
+        return np.concatenate(arrs, **kw).view(cls)
+    
+    def remove(self, other, **kw):
+        return np.setdiff1d(self, other, assume_unique=True, **kw).view(self.__class__)
+    
+    def intersect(self, other, **kw):
+        return np.intersect1d(self, other, assume_unique=True, **kw).view(self.__class__)
+    
+
+class BooleanOperationError(NotImplementedError):
+    """ Raised when a logical operation is performed on a non-logical array """
+    def __init__(self, arr):
+        msg = f'Logical operations are only valid on Boolean arrays, not {arr.dtype}'
+        super().__init__(msg)
\ No newline at end of file
diff --git a/starsim/utils.py b/starsim/utils.py
index 89f25273..ed3be7db 100644
--- a/starsim/utils.py
+++ b/starsim/utils.py
@@ -12,7 +12,7 @@
 # %% Helper functions
 
 # What functions are externally visible -- note, this gets populated in each section below
-__all__ = ['ndict', 'omerge', 'omergeleft', 'warn', 'unique', 'find_contacts', 'get_subclasses', 'all_subclasses']
+__all__ = ['ndict', 'dictmerge', 'dictmergeleft', 'warn', 'unique', 'find_contacts', 'get_subclasses', 'all_subclasses']
 
 
 class ndict(sc.objdict):
@@ -125,11 +125,12 @@ def __iadd__(self, dict2):
         return self
 
 
-def omerge(*args, **kwargs):
+def dictmerge(*args, **kwargs):
     """ Merge things into an objdict, using standard order """
     return sc.objdict(sc.mergedicts(*args, **kwargs))
 
-def omergeleft(*args, **kwargs):
+
+def dictmergeleft(*args, **kwargs):
     """ Merge things into an odict, using opposite order to allow defaults to be supplied second """
     if len(args) == 1 and len(kwargs):
         new = args[0]
@@ -138,10 +139,11 @@ def omergeleft(*args, **kwargs):
         new = args[0]
         default = args[1]
     else:
-        errormsg = 'Expecting either two arguments, or one argument and kwargs; for any other arrangement, use ss.omerge()'
+        errormsg = 'Expecting either two arguments, or one argument and kwargs; for any other arrangement, use ss.dictmerge()'
         raise ValueError(errormsg)
     return sc.objdict(sc.mergedicts(default, new))
 
+
 def warn(msg, category=None, verbose=None, die=None):
     """ Helper function to handle warnings -- shortcut to warnings.warn """
 
@@ -171,6 +173,7 @@ def warn(msg, category=None, verbose=None, die=None):
 
     return
 
+
 def unique(arr):
     """
     Find the unique elements and counts in an array.
@@ -182,6 +185,7 @@ def unique(arr):
     counts = counts[unique]
     return unique, counts
 
+
 def find_contacts(p1, p2, inds):  # pragma: no cover
     """
     Variation on Network.find_contacts() that avoids sorting.
@@ -199,6 +203,7 @@ def find_contacts(p1, p2, inds):  # pragma: no cover
             pairing_partners.add(p1[i])
     return pairing_partners
 
+
 def get_subclasses(cls):
     for subclass in cls.__subclasses__():
         yield from get_subclasses(subclass)
@@ -241,94 +246,12 @@ def set_seed_regular(seed):
 
     return
 
-# %% Simple array operations
-
-__all__ += ['true', 'false', 'defined', 'undefined']
-
-@nb.njit(cache=True)
-def _true(uids, values):
-    """
-    Returns the UIDs for indices where the value evaluates as True
-    """
-    out = np.empty(len(uids), dtype=uids.dtype)
-    j = 0
-    for i in range(len(values)):
-        out[j] = uids[i]
-        if values[i]:
-            j += 1
-    out = out[0:j]
-    return out
-
-@nb.njit(cache=True)
-def _false(uids, values):
-    """
-    Returns the UIDs for indices where the value evaluates as False
-    """
-    out = np.empty(len(uids), dtype=uids.dtype)
-    j = 0
-    for i in range(len(values)):
-        out[j] = uids[i]
-        if not values[i]:
-            j += 1
-    out = out[0:j]
-    return out
-
-def true(state):
-    """
-    Returns the UIDs of the values of the array that are true
-
-    Args:
-        state (State, UIDArray)
-
-    **Example**::
-
-        inds = ss.true(people.alive) # Returns array of UIDs of alive agents
-    """
-    return _true(state.uid.__array__(), state.__array__())
-
-def false(state):
-    """
-    Returns the indices of the values of the array that are false.
-
-    Args:
-        state (State, UIDArray)
-
-    **Example**::
-
-        inds = ss.false(people.alive) # Returns array of UIDs of dead agents
-    """
-    return _false(state.uid.__array__(), state.__array__())
-
-def defined(arr):
-    """
-    Returns the indices of the values of the array that are not-nan.
-
-    Args:
-        arr (array): any array
-
-    **Example**::
-
-        inds = ss.defined(np.array([1,np.nan,0,np.nan,1,0,1]))
-    """
-    return (~np.isnan(arr)).nonzero()[-1]
-
-def undefined(arr):
-    """
-    Returns the indices of the values of the array that are not-nan.
-
-    Args:
-        arr (array): any array
-
-    **Example**::
-
-        inds = ss.defined(np.array([1,np.nan,0,np.nan,1,0,1]))
-    """
-    return np.isnan(arr).nonzero()[-1]
 
 # %% Data cleaning and processing
 
 __all__ += ['standardize_data']
 
+
 def standardize_data(data=None, metadata=None, max_age=120, min_year=1800):
     """
     Args:
diff --git a/starsim/version.py b/starsim/version.py
index 6f43b155..2878d0f2 100644
--- a/starsim/version.py
+++ b/starsim/version.py
@@ -4,6 +4,6 @@
 
 __all__ = ['__version__', '__versiondate__', '__license__']
 
-__version__ = '0.3.4'
-__versiondate__ = '2024-04-18'
+__version__ = '0.4.0'
+__versiondate__ = '2024-04-24'
 __license__ = f'Starsim {__version__} ({__versiondate__}) — © 2023-2024 by IDM'


Test patch:
diff --git a/tests/baseline.json b/tests/baseline.json
index 42067cfa..3df135e1 100644
--- a/tests/baseline.json
+++ b/tests/baseline.json
@@ -1,17 +1,17 @@
 {
   "summary": {
-    "n_alive": 9648.460784313726,
-    "new_deaths": 7.549019607843137,
-    "cum_deaths": 760.0,
+    "n_alive": 9647.450980392157,
+    "new_deaths": 7.490196078431373,
+    "cum_deaths": 755.0,
     "pregnancy_pregnancies": 0.0,
     "pregnancy_births": 0.0,
     "pregnancy_cbr": 0.0,
     "hiv_n_on_art": 0.0,
-    "hiv_n_susceptible": 8781.696078431372,
-    "hiv_n_infected": 866.7647058823529,
-    "hiv_prevalence": 0.08997817294776833,
-    "hiv_new_infections": 16.84313725490196,
-    "hiv_cum_infections": 1718.0,
-    "hiv_new_deaths": 7.549019607843137
+    "hiv_n_susceptible": 8785.009803921568,
+    "hiv_n_infected": 862.4411764705883,
+    "hiv_prevalence": 0.08951620561257022,
+    "hiv_new_infections": 16.705882352941178,
+    "hiv_cum_infections": 1704.0,
+    "hiv_new_deaths": 7.490196078431373
   }
 }
\ No newline at end of file
diff --git a/tests/benchmark.json b/tests/benchmark.json
index a5895ba8..2c18591e 100644
--- a/tests/benchmark.json
+++ b/tests/benchmark.json
@@ -2,12 +2,12 @@
   "time": {
     "people": 0.001,
     "initialize": 0.023,
-    "run": 0.31
+    "run": 0.355
   },
   "parameters": {
     "n_agents": 10000,
     "n_years": 20,
     "dt": 0.2
   },
-  "cpu_performance": 0.9635620582658682
+  "cpu_performance": 0.974601487661388
 }
\ No newline at end of file
diff --git a/tests/test_base.py b/tests/test_base.py
index 25638f2c..67b3467b 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -20,8 +20,11 @@ def test_people():
     del ppl
 
     # Possible to initialize people with extra states, e.g. a geolocation
+    def geo_func(n):
+        locs = [1,2,3]
+        return np.random.choice(locs, n)
     extra_states = [
-        ss.State('geolocation', int, lambda n: np.random.choice([1,2,3],n)),
+        ss.FloatArr('geolocation', default=geo_func),
     ]
     ppl = ss.People(100, extra_states=extra_states)
 
@@ -54,7 +57,7 @@ def test_networks():
 
 
 def test_microsim(do_plot=False):
-    sc.heading('Test making people and providing them to a sim')
+    sc.heading('Test small HIV simulation')
 
     # Make HIV module
     hiv = ss.HIV()
@@ -81,6 +84,7 @@ def test_microsim(do_plot=False):
 
 
 def test_ppl_construction():
+    sc.heading('Test making people and providing them to a sim')
 
     def init_debut(module, sim, uids):
         # Test setting the mean debut age by sex, 16 for men and 21 for women.
@@ -105,6 +109,32 @@ def init_debut(module, sim, uids):
     return sim
 
 
+def test_arrs():
+    sc.heading('Testing Arr objects')
+    o = sc.objdict()
+    
+    # Create a sim with only births
+    pars = dict(n_agents=1000, diseases='sis', networks='random')
+    p1 = sc.mergedicts(pars, birth_rate=10)
+    p2 = sc.mergedicts(pars, death_rate=10)
+    s1 = ss.Sim(pars=p1).run()
+    s2 = ss.Sim(pars=p2).run()
+    
+    # Tests
+    assert len(s1.people.auids) > len(s2.people.auids), 'Should have more people with births'
+    assert np.array_equal(s1.people.age, s1.people.age.raw[s1.people.auids]), 'Different ways of indexing age should match'
+    assert np.array_equal(s1.people.alive.uids, s1.people.auids), 'Active/alive agents should match'
+    assert np.all(s2.people.alive), 'All agents should be alive when indexed like this'
+    assert not np.all(s2.people.alive.raw), 'Some agents should not be alive when indexed like this'
+    assert np.array_equal(~s1.people.female, s1.people.male), 'Definition of men does not match'
+    assert isinstance(s1.people.age < 5, ss.BoolArr), 'Performing logical operations should return a BoolArr'
+    
+    o.s1 = s1
+    o.s2 = s2   
+    
+    return o
+
+
 
 # %% Run as a script
 if __name__ == '__main__':
@@ -119,6 +149,7 @@ def init_debut(module, sim, uids):
     nw1, nw2, nw3 = test_networks()
     sim1 = test_microsim(do_plot)
     sim2 = test_ppl_construction()
+    sims = test_arrs()
 
     sc.toc(T)
     pl.show()
diff --git a/tests/test_dists.py b/tests/test_dists.py
index ff65a66e..10599299 100644
--- a/tests/test_dists.py
+++ b/tests/test_dists.py
@@ -127,7 +127,7 @@ def apply(self, sim):
             # Create an extra agent
             preg = ss.Pregnancy(rel_fertility=0) # Ensure no default births
             preg.initialize(sim)
-            new_uids = np.array([len(sim.people)]) # Hack since make_embryos doesn't return UIDs
+            new_uids = ss.uids([len(sim.people)]) # Hack since make_embryos doesn't return UIDs
             preg.make_embryos(sim, np.array([0])) # Assign 0th agent to be the "mother"
             assert len(new_uids) == 1
             sim.people.age[new_uids] = -100 # Set to a very low number to never reach debut age
diff --git a/tests/test_syphilis.py b/tests/test_syphilis.py
index 5a8341fd..d9e5855e 100644
--- a/tests/test_syphilis.py
+++ b/tests/test_syphilis.py
@@ -140,7 +140,7 @@ def test_syph_intvs(dt=1, n_agents=500, do_plot=False):
         label='syph_screening',
     )
 
-    treat_eligible = lambda sim: sim.get_intervention('syph_screening').outcomes['positive']
+    treat_eligible = lambda sim: ss.uids(sim.get_intervention('syph_screening').outcomes['positive'])
     bpg = ss.syph_treatment(
         prob=0.9,
         product='bpg',