add the ability to reassemble live calculated TimeChunked quantities

docs/histogram_properties.md

@@ -108,6 +108,26 @@ Note that it is a _very_ bad idea to change this after you have already written
 for a simulation. You will almost certainly end up getting inconsistent results. Always set `histogram_delta_t_Gyr`
 before your first `tangos write`.
 
+Reassembling live calculations
+------------------------------
+
+It is possible to define a [live calculation](live_calculation.md) that manipulates an already stored histogram
+property (like `SFR_histogram`) and returns a new histogram. Because such a calculation only requires data already stored
+in the database, it can be performed on the fly and reassembled accordingly in the same way as the original stored
+histogram data.
+
+The property `SpecSFR_histogram`, for example, calculates the specific star formation rate history of a halo given `SFR_histogram` and
+normalizing it by the mass at every intermediate step. We can call this in the same way as any other histogram property, but
+this time we are doing a live calculation, which is called like any other live calculation (in this case there are
+no relevant arguments to give, but any argument taken by the live calculation will work like normal).
+
+```
+halo.calculate('reassemble(SpecSFR_histogram())')
+```
+
+It is still important to make sure that the time resolution is set correctly. For any live calculation acting on a histogram
+property, the time resolution, defined for each simulation by `sim["histogram_delta_t_Gyr"]`, should be set to what it 
+was when the stored histogram data was originally calculated.
 
 Exercising caution
 -------------------

tangos/core/extraction_patterns.py

@@ -131,7 +131,7 @@ def _postprocess_one_result(self, property_object):
 
         if hasattr(self._providing_class, 'reassemble'):
             instance = self._providing_class(property_object.halo.timestep.simulation)
-            return instance.reassemble(property_object, *self._options)
+            return instance.reassemble(property_object, property_object.halo, *self._options)
         else:
             self._setup_data_mapper(property_object)
             return self._mapper.get(property_object)

tangos/live_calculation/__init__.py

@@ -366,6 +366,8 @@ def __init__(self, *tokens):
         super(LiveProperty, self).__init__()
         self._name = str(tokens[0])
         self._inputs = list(tokens[1:])
+        self._evaluation_with_reassemble = True
+        self._evaluation_options= []
 
     def __str__(self):
         return self._name + "(" + (",".join(str(x) for x in self._inputs)) + ")"
@@ -392,6 +394,18 @@ def _calculation_retrieves(self):
         result = result.union(providing_instance.requires_property())
         return result
 
+    def set_reassemble(self):
+        self._evaluation_with_reassemble = True
+
+    def set_raw(self):
+        self._evaluation_with_reassemble = False
+
+    def set_evaluation_options(self,*options):
+        self._evaluation_options = options
+
+    def _evaluate(self, halos, input_descriptions, input_values):
+        return self._evaluate_function(halos, input_descriptions, input_values, *self._evaluation_options)
+
     def values_and_description(self, halos):
         if len(halos)==0:
             return np.array([[]]), None
@@ -402,7 +416,7 @@ def values_and_description(self, halos):
             input_values.append(iv)
             input_descriptions.append(id)
 
-        calculator, results = self._evaluate_function(halos, input_descriptions, input_values)
+        calculator, results = self._evaluate(halos, input_descriptions, input_values)
 
         return results, calculator
 
@@ -413,14 +427,17 @@ def _input_value_and_description(self, input_id, halos):
             raise ValueError("Functions cannot receive more than one value per input")
         return val[0], desc
 
-    def _evaluate_function(self, halos, input_descriptions, input_values):
+    def _evaluate_function(self, halos, input_descriptions, input_values, *options):
         from .. import properties
         sim = consistent_collection.consistent_simulation_from_halos(halos)
         results = []
         calculator = properties.providing_class(self.name())(sim, *input_descriptions)
         for inputs in zip(halos, *input_values):
             if self._has_required_properties(inputs[0]) and all([x is not None for x in inputs]):
-                results.append(calculator.live_calculate_named(self.name(), *inputs))
+                if hasattr(calculator,'reassemble') & self._evaluation_with_reassemble is True:
+                    results.append(calculator.reassemble(self, inputs[0], *options))
+                else:
+                    results.append(calculator.live_calculate_named(self.name(), *inputs))
             else:
                 results.append(None)
         return calculator, self._as_1xn_array(results)
@@ -488,6 +505,7 @@ def __init__(self, *tokens):
         super(BuiltinFunction, self).__init__(*tokens)
         self._func = self.__registered_functions[self._name]['function']
         self._info = self.__registered_functions[self._name]
+        self.set_raw()
         for i in range(len(self._inputs)):
             assert_class = self._get_input_option(i, 'assert_class')
             if assert_class is not None and not isinstance(self._inputs[i], assert_class):

tangos/live_calculation/builtin_functions/reassembly.py

@@ -1,23 +1,24 @@
 from __future__ import absolute_import
 from tangos.core import extraction_patterns
 from . import BuiltinFunction
-from .. import StoredProperty, FixedInput
+from .. import StoredProperty, FixedInput, LiveProperty
 
 
 @BuiltinFunction.register
 def raw(halos, values):
     return values
-raw.set_input_options(0, assert_class=StoredProperty)
 
 @raw.set_initialisation
 def raw_initialisation(input):
-    input.set_extraction_pattern(extraction_patterns.HaloPropertyRawValueGetter())
+    if isinstance(input, LiveProperty):
+        input.set_raw()
+    else:
+        input.set_extraction_pattern(extraction_patterns.HaloPropertyRawValueGetter())
 
 
 @BuiltinFunction.register
 def reassemble(halos, values, *options):
     return values
-reassemble.set_input_options(0, assert_class=StoredProperty)
 
 @reassemble.set_initialisation
 def reassemble_initialisation(input, *options):
@@ -27,6 +28,9 @@ def reassemble_initialisation(input, *options):
             options_values.append(option.proxy_value())
         else:
             raise TypeError("Options to 'reassemble' must be fixed numbers or strings")
-
-    input.set_extraction_pattern(
-        extraction_patterns.HaloPropertyValueWithReassemblyOptionsGetter(*options_values))
+    if isinstance(input,LiveProperty):
+        input.set_reassemble()
+        input.set_evaluation_options(*options_values)
+    else:
+        input.set_extraction_pattern(
+            extraction_patterns.HaloPropertyValueWithReassemblyOptionsGetter(*options_values))

tangos/properties/__init__.py

@@ -5,6 +5,7 @@
 from six.moves import zip
 import importlib
 import warnings
+from ..live_calculation import LiveProperty
 import functools
 from .. import input_handlers
 from .. import util
@@ -229,10 +230,10 @@ class TimeChunkedProperty(PropertyCalculation):
     minimum_store_Gyr = 0.5
 
     def __init__(self, simulation):
-        self.pixel_delta_t_Gyr = simulation.get("histogram_delta_t_Gyr", self.pixel_delta_t_Gyr)
+        if simulation is not None:
+            self.pixel_delta_t_Gyr = simulation.get("histogram_delta_t_Gyr", self.pixel_delta_t_Gyr)
         super(TimeChunkedProperty, self).__init__(simulation)
 
-
     def bin_index(self, time):
         """Convert a time (Gyr) to a bin index in the histogram"""
         index = int(time/self.pixel_delta_t_Gyr)
@@ -245,7 +246,8 @@ def store_slice(self, time):
         they should store."""
         return slice(self.bin_index(time-self.minimum_store_Gyr), self.bin_index(time))
 
-    def reassemble(self, property, reassembly_type='major'):
+
+    def reassemble(self, property, halo, reassembly_type='major'):
         """Reassemble a histogram by suitable treatment of the merger tree leading up to the current halo.
 
         This function is normally called by the framework (see Halo.get_data_with_reassembly_options) and you would
@@ -266,16 +268,24 @@ def reassemble(self, property, reassembly_type='major'):
         from tangos import relation_finding as rfs
 
         if reassembly_type=='major':
-            return self._reassemble_using_finding_strategy(property, strategy = rfs.MultiHopMajorProgenitorsStrategy)
+            return self._reassemble_using_finding_strategy(property, halo, strategy = rfs.MultiHopMajorProgenitorsStrategy)
         elif reassembly_type=='major_across_simulations':
-            return self._reassemble_using_finding_strategy(property, strategy = rfs.MultiHopMajorProgenitorsStrategy,
+            return self._reassemble_using_finding_strategy(property, halo, strategy = rfs.MultiHopMajorProgenitorsStrategy,
                                                            strategy_kwargs = {'target': None, 'one_simulation': False})
         elif reassembly_type=='sum':
-            return self._reassemble_using_finding_strategy(property, strategy = rfs.MultiHopAllProgenitorsStrategy)
+            return self._reassemble_using_finding_strategy(property, halo, strategy = rfs.MultiHopAllProgenitorsStrategy)
         elif reassembly_type=='place':
-            return self._place_data(property.halo.timestep.time_gyr, property.data_raw)
+            if not isinstance(property, LiveProperty):
+                place_data = property.data_raw
+            else:
+                place_data = halo.calculate(property.__str__())
+            return self._place_data(property.halo.timestep.time_gyr, place_data)
+
         elif reassembly_type=='raw':
-            return property.data_raw
+            if not isinstance(property, LiveProperty):
+                return property.data_raw
+            else:
+                return halo.calculate('raw('+property.__str__()+')')
         else:
             raise ValueError("Unknown reassembly type")
 
@@ -286,9 +296,11 @@ def _place_data(self, time, raw_data):
         final[start:] = raw_data
         return final
 
-    def _reassemble_using_finding_strategy(self, property, strategy, strategy_kwargs={}):
-        name = property.name.text
-        halo = property.halo
+    def _reassemble_using_finding_strategy(self, property, halo, strategy, strategy_kwargs={}):
+        if not isinstance(property, LiveProperty):
+            name = property.name.text
+        else:
+            name = property.__str__()
         t, stack = halo.calculate_for_descendants("t()", "raw(" + name + ")", strategy=strategy, strategy_kwargs=strategy_kwargs)
         final = np.zeros(self.bin_index(t[0]))
         previous_time = -1

tangos/properties/pynbody/SF.py

@@ -1,4 +1,4 @@
-from .. import TimeChunkedProperty
+from .. import TimeChunkedProperty, LiveHaloProperties
 from . import pynbody_handler_module, PynbodyPropertyCalculation
 
 import numpy as np
@@ -33,10 +33,36 @@ def calculate(self, halo, existing_properties):
 
         return M
 
-    def reassemble(self, *options):
-        reassembled = super(StarFormHistogram, self).reassemble(*options)
+    def reassemble(self, halo, *options):
+        reassembled = super(StarFormHistogram, self).reassemble(halo, *options)
         return reassembled/1e9 # Msol per Gyr -> Msol per yr
 
+class SpecStarFormationHistogram(TimeChunkedProperty,LiveHaloProperties):
+    names = "SpecSFR_histogram"
+
+    def __init__(self, simulation):
+        super(SpecStarFormationHistogram, self).__init__(simulation)
+
+    def plot_xlabel(self):
+        return "t/Gyr"
+
+    def plot_ylabel(self):
+        return r"$SFR/M_{\odot} yr^{-1}$"
+
+    def requires_property(self):
+        return ['Mstar','SFR_histogram']
+
+    def live_calculate(self, halo, *args):
+        sfr = halo.calculate('raw(SFR_histogram)')
+        t_sfr = halo.timestep.time_gyr - np.arange(len(sfr))*self.pixel_delta_t_Gyr
+        Marray = halo['Mstar'] - np.cumsum(sfr*self.pixel_delta_t_Gyr)
+        Marray = Marray[::-1]
+        return sfr/Marray
+
+    def reassemble(self, halo, *options):
+        reassembled = super(SpecStarFormationHistogram, self).reassemble(halo, *options)
+        return reassembled / 1e9  # Gyr^-1 -> yr^-1
+
 class StarForm(PynbodyPropertyCalculation):
     names = "SFR_10Myr", "SFR_100Myr"
 

tangos/properties/pynbody/profile.py

@@ -21,7 +21,6 @@ def plot_ylabel(self):
 
     def _get_profile(self, halo, maxrad):
         import pynbody
-
         delta = self.plot_xdelta()
         nbins = int(maxrad / delta)
         maxrad = delta * (nbins + 1)
@@ -50,6 +49,7 @@ class BaryonicHaloDensityProfile(HaloDensityProfile):
 
     @centred_calculation
     def calculate(self, data, existing_properties):
+        import pynbody
         gas_a, gas_b = self._get_profile(data.gas, existing_properties["max_radius"])
-        star_a, star_b = self._get_profile(data.star, existing_properties["max_radius"])
+        star_a, star_b = self._get_profile(data.star[pynbody.filt.HighPass('tform',0)], existing_properties["max_radius"])
         return gas_a, gas_b, star_a, star_b

tests/test_live_calculation.py

@@ -81,12 +81,25 @@ class DummyPropertyWithReassemblyOptions(properties.PropertyCalculation):
     names = "dummy_property_with_reassembly"
 
     @classmethod
-    def reassemble(cls, property, test_option=25):
+    def reassemble(cls, property, halo, test_option=25):
         if test_option=='actual_data':
             return property.data_raw
         else:
             return test_option
 
+class DummyLivePropertyWithReassembly(properties.LivePropertyCalculation):
+    names = "dummy_live_property_with_reassembly"
+
+    def live_calculate(self, db_halo_entry, factor=2):
+        return db_halo_entry.calculate('raw(dummy_property_with_reassembly)') * factor
+
+    @classmethod
+    def reassemble(cls, property, halo, test_option=10):
+        if test_option == 'actual_data':
+            return halo.calculate('raw('+property.__str__()+')')
+        else:
+            return halo.calculate('raw('+property.__str__()+')') * test_option
+
 class LivePropertyRequiringRedirectedProperty(properties.LivePropertyCalculation):
     names = "first_BHs_BH_mass"
 
@@ -251,6 +264,27 @@ def test_reassembly():
     # our dummy reassembly function can return the actual value if suitably requested
     assert h.calculate('reassemble(dummy_property_with_reassembly, "actual_data")') == 101
 
+def test_live_reassembly():
+    h = tangos.get_halo("sim/ts1/1")
+    h['dummy_property_with_reassembly'] = 101
+
+    # default reassembly takes in the dummy property and multiplies it first by 2, then by 10
+    assert h.calculate('dummy_live_property_with_reassembly()') == 2020
+
+    # default live_calculation without reassemble takes dummy property and multiplies by 2
+    assert h.calculate('raw(dummy_live_property_with_reassembly())') == 202
+
+    # ability to give input into live calculation
+    assert h.calculate('raw(dummy_live_property_with_reassembly(4))') == 404
+
+    # pass option to reassembly and change how the dummpy property is manipulated
+    assert h.calculate('reassemble(dummy_live_property_with_reassembly(),3)') == 606
+
+    # reassemble on live calculation with new variable
+    assert h.calculate('reassemble(dummy_live_property_with_reassembly(3))') == 3030
+
+    # will return the original value of the live calculation without reassembly if requested correctly
+    assert h.calculate('reassemble(dummy_live_property_with_reassembly(3),"actual_data")') == 303
 
 def test_liveproperty_requiring_redirection():
     h = tangos.get_halo("sim/ts1/1")