Merge pull request #30 from ClimateImpactLab/fusion

WIP: Merge Release-2.0 changes into master

Author: jrising

openest/curves/smart_linextrap.py

@@ -60,11 +60,12 @@ def __call__(self, ds):
         values = self.curve(ds)
         indeps = np.stack((ds[indepvar] for indepvar in self.indepvars), axis=-1)
 
-        return linextrap.replace_oob(values, indeps, self.curve.get_univariate(), self.bounds, self.margins, self.scaling)
+        return linextrap.replace_oob(values, indeps, self.curve.univariate, self.bounds, self.margins, self.scaling)
 
-    def get_univariate(self):
+    @property
+    def univariate(self):
         """Return a UnivariateCurve version of this curve."""
-        return linextrap.LinearExtrapolationCurve(self.curve.get_univariate(), self.bounds, self.margins,
+        return linextrap.LinearExtrapolationCurve(self.curve.univariate, self.bounds, self.margins,
                                                   self.scaling, lambda xs: xs[:, 0])
 
 

openest/curves/ushape_numeric.py

@@ -1,4 +1,6 @@
 import numpy as np
+import xarray as xr
+from openest.generate import fast_dataset
 from openest.curves.basic import UnivariateCurve
 
 
@@ -33,11 +35,10 @@ def __init__(self, curve, midtemp, gettas, ordered='maintain', fillxxs=None, fil
 
     def __call__(self, xs):
         values = self.curve(xs)
-        tas = np.array(self.gettas(xs))
+        tas = tas_original = np.array(self.gettas(xs))
 
         # Add in the grid for completeness
         if len(self.fillxxs) > 0:
-            tas_saved = tas
             tas = np.concatenate((tas, self.fillxxs))
             values = np.concatenate((values, self.fillyys))
 
@@ -65,10 +66,10 @@ def __call__(self, xs):
             tasorder = np.empty(len(increasing))
             tasorder[order] = increasing
             # Return just the given values
-            return tasorder[:len(xs)]
+            return tasorder[:len(tas_original)]
         else:
             if len(self.fillxxs) > 0:
-                tokeeps = order < len(tas_saved)
+                tokeeps = order < len(tas_original)
                 lowkeeps = tokeeps[orderedtas < self.midtemp]
                 highkeeps = tokeeps[orderedtas >= self.midtemp]
 
@@ -157,7 +158,10 @@ def __call__(self, xs):
             highindicesofordered = np.maximum.accumulate(highindicesofordered)
 
         # Construct the results
-        if len(xs.shape) == 2:
+        if isinstance(xs, xr.Dataset):
+            newxs = fast_dataset.reorder_coord(xs, 'time', np.concatenate((order[lowindicesofordered], order[highindicesofordered])))
+            increasingresults = self.tmarginal_curve(newxs)  # ordered low..., high...
+        elif len(xs.shape) == 2:
             increasingresults = np.concatenate((self.tmarginal_curve(xs[order[lowindicesofordered], :]),
                                                 self.tmarginal_curve(xs[order[highindicesofordered], :])))  # ordered low..., high...
         else:

openest/generate/fast_dataset.py

@@ -146,6 +146,9 @@ def isel(self, **kwargs):
 
         return FastDataset(newdata_vars, newcoords, self.attrs)
 
+    def load(self):
+        pass # always ready
+    
     def __getitem__(self, name):
         if name not in self._variables:
             if name == 'time.year' and 'time' in self._variables:
@@ -472,6 +475,40 @@ def assert_index_equal(one, two):
             assert np.array_equal(two, one), "Not equal: %s <> %s" % (str(two), str(one))
             return one
 
+def reorder_coord(ds, dim, indices):
+    assert dim in ['time', 'region'] # assume a time x region ds
+    assert dim in ds.coords
+    
+    newvars = {}
+    for var in ds.variables:
+        try:
+            dims = ds[var].coords
+            if len(dims) == 1:
+                if dim not in dims:
+                    newvars[var] = ds[var]
+                else:
+                    newvars[var] = ([dim], ds[var].values[indices])
+            elif len(dims) == 2 and dim in dims:
+                if dims.index(dim) == 0:
+                    newvars[var] = (dims, ds[var].values[:, indices])
+                else:
+                    newvars[var] = (dims, ds[var].values[indices, :])
+            else:
+                tindex = list(dims).index(dim)
+                allindices = [slice(None)] * len(dims)
+                allindices[tindex] = indices
+                newvars[var] = (dims, ds[var].values[tuple(allindices)])
+        except Exception:
+            print("Failed to reorder %s for %s" % (var, ds))
+            raise
+
+    coords = ds.coords
+    coords[dim] = ds[dim][indices]
+    newds = ds.__class__(newvars, coords=coords) # work for Dataset or FastDataset
+    newds.load()
+
+    return newds
+    
 FastDataArray.__array_priority__ = 80
 xr.core.ops.inject_binary_ops(FastDataArray)
 

openest/generate/functions.py

@@ -868,11 +868,15 @@ class AuxiliaryResult(calculation.Calculation):
     """
     Produce an additional output, but then pass the main result on.
     """
-    def __init__(self, subcalc_main, subcalc_aux, auxname):
-        super(AuxiliaryResult, self).__init__([subcalc_main.unitses[0], subcalc_aux.unitses[0]] + subcalc_main.unitses[1:])
+    def __init__(self, subcalc_main, subcalc_aux, auxname, keeplastonly=True):
+        if keeplastonly:
+            super(AuxiliaryResult, self).__init__([subcalc_main.unitses[0], subcalc_aux.unitses[0]] + subcalc_main.unitses[1:])
+        else:
+            super(AuxiliaryResult, self).__init__([subcalc_main.unitses[0]] + subcalc_aux.unitses + subcalc_main.unitses[1:])
         self.subcalc_main = subcalc_main
         self.subcalc_aux = subcalc_aux
         self.auxname = auxname
+        self.keeplastonly = keeplastonly
 
     def format(self, lang, *args, **kwargs):
         beforeauxlen = len(formatting.format_labels)
@@ -884,7 +888,7 @@ def format(self, lang, *args, **kwargs):
     def apply(self, region, *args, **kwargs):
         subapp_main = self.subcalc_main.apply(region, *args, **kwargs)
         subapp_aux = self.subcalc_aux.apply(region, *args, **kwargs)
-        return AuxiliaryResultApplication(region, subapp_main, subapp_aux)
+        return AuxiliaryResultApplication(region, subapp_main, subapp_aux, self.keeplastonly)
 
     def partial_derivative(self, covariate, covarunit):
         """
@@ -899,7 +903,10 @@ def column_info(self):
         infos_aux = self.subcalc_aux.column_info()
         infos_aux[0]['name'] = self.auxname
 
-        return [infos_main[0], infos_aux[0]] + infos_main[1:]
+        if self.keeplastonly:
+            return [infos_main[0], infos_aux[0]] + infos_main[1:]
+        else:
+            return [infos_main[0]] + infos_aux + infos_main[1:]
 
     @staticmethod
     def describe():
@@ -927,16 +934,20 @@ class AuxiliaryResultApplication(calculation.Application):
     """
     Perform both main and auxiliary calculation, and order as main[0], aux, main[1:]
     """
-    def __init__(self, region, subapp_main, subapp_aux):
+    def __init__(self, region, subapp_main, subapp_aux, keeplastonly):
         super(AuxiliaryResultApplication, self).__init__(region)
         self.subapp_main = subapp_main
         self.subapp_aux = subapp_aux
+        self.keeplastonly = keeplastonly
 
     def push(self, ds):
         for yearresult in self.subapp_main.push(ds):
             for yearresult_aux in self.subapp_aux.push(ds):
-                next # Just take the last one
-            yield list(yearresult[0:2]) + [yearresult_aux[1]] + list(yearresult[2:])
+                pass # Just take the last one
+            if self.keeplastonly:
+                yield list(yearresult[0:2]) + [yearresult_aux[1]] + list(yearresult[2:])
+            else:
+                yield list(yearresult[0:2]) + list(yearresult_aux[1:]) + list(yearresult[2:])
 
     def done(self):
         self.subapp_aux.done()

openest/generate/smart_curve.py

@@ -22,8 +22,9 @@ def __init__(self):
     def __call__(self, ds):
         raise NotImplementedError("call not implemented")
 
-    def get_univariate(self):
-        raise NotImplementedError("get_univariate not implemented")
+    @property
+    def univariate(self):
+        raise NotImplementedError("univariate not implemented")
 
     def format(self, lang):
         raise NotImplementedError()
@@ -120,7 +121,8 @@ def __call__(self, ds):
 
         return result
 
-    def get_univariate(self):
+    @property
+    def univariate(self):
         return curve.ZeroInterceptPolynomialCurve([-np.inf, np.inf], self.coeffs)
 
     def format(self, lang):
@@ -171,6 +173,7 @@ class SumByTimePolynomialCurve(SmartCurve):
     def __init__(self, coeffmat, variables, allow_raising=False, descriptions=None):
         super(SumByTimePolynomialCurve, self).__init__()
         self.coeffmat = coeffmat # K x T
+        assert len(self.coeffmat.shape) == 2
         self.variables = variables
         self.allow_raising = allow_raising
         if descriptions is None:
@@ -259,7 +262,8 @@ def __call__(self, ds):
 
             raise ex
 
-    def get_univariate(self):
+    @property
+    def univariate(self):
         return curve.CubicSplineCurve(self.knots, self.coeffs)
 
 class TransformCoefficientsCurve(SmartCurve):
@@ -351,22 +355,30 @@ def __init__(self, curve, offset):
         self.offset = offset
 
     def __call__(self, ds):
-        return self.curve(ds) - self.offset
+        return self.curve(ds) + self.offset
 
-    def get_univariate(self):
-        return curve.ShiftedCurve(self.curve.get_univariate(), self.offset)
+    @property
+    def univariate(self):
+        return curve.ShiftedCurve(self.curve.univariate, self.offset)
 
     def format(self, lang):
-        return formatting.build_recursive({'latex': r"(%s - " + str(self.offset) + ")",
-                                           'julia': r"(%s - " + str(self.offset) + ")"},
+        return formatting.build_recursive({'latex': r"(%s + " + str(self.offset) + ")",
+                                           'julia': r"(%s + " + str(self.offset) + ")"},
                                           lang, self.curve)
 
 class ClippedCurve(curve.ClippedCurve, SmartCurve):
-    def get_univariate(self):
-        return curve.ClippedCurve(self.curve.get_univariate(), self.cliplow)
+    @property
+    def univariate(self):
+        return curve.ClippedCurve(self.curve.univariate, self.cliplow)
+
+class OtherClippedCurve(curve.OtherClippedCurve, SmartCurve):
+    @property
+    def univariate(self):
+        return curve.OtherClippedCurve(self.clipping_curve.univariate, self.curve.univariate, self.clipy)
 
 class MinimumCurve(curve.MinimumCurve, SmartCurve):
-    def get_univariate(self):
-        return curve.MinimumCurve(self.curve1.get_univariate(), self.curve2.get_univariate())
+    @property
+    def univariate(self):
+        return curve.MinimumCurve(self.curve1.univariate, self.curve2.univariate)
 
 

openest/generate/yearly.py

@@ -57,22 +57,26 @@ def describe():
                     description="Apply the results of a previous calculation to a curve.")
 
 class YearlyCoefficients(Calculation):
-    def __init__(self, units, curvegen, curve_description, getter=lambda curve: curve.yy, weather_change=lambda region, x: x):
+    def __init__(self, units, curvegen, curve_description, getter=lambda curve: curve.yy, weather_change=lambda region, x: x, label='response'):
         super(YearlyCoefficients, self).__init__([units])
         assert isinstance(curvegen, CurveGenerator)
 
         self.curvegen = curvegen
         self.curve_description = curve_description
         self.getter = getter
         self.weather_change = weather_change
+        self.label = label
 
     def apply(self, region, *args):
         def generate(region, year, temps, **kw):
             curve = self.curvegen.get_curve(region, year, *args, weather=temps) # Passing in original (not weather-changed) data
 
             coeffs = self.getter(region, year, temps, curve)
-            if len(temps) == len(coeffs):
-                result = np.sum(self.weather_change(region, temps).dot(coeffs))
+            temps2 = self.weather_change(region, temps)
+            if np.isscalar(temps2) and np.isscalar(coeffs):
+                result = temps2 * coeffs
+            elif len(temps) == len(coeffs):
+                result = np.sum(np.array(temps).dot(coeffs))
             else:
                 raise RuntimeError("Unknown format for temps: " + str(temps.shape) + " <> len " + str(coeffs))
 
@@ -87,7 +91,7 @@ def generate(region, year, temps, **kw):
 
     def column_info(self):
         description = "The combined result of yearly values, with coefficients from %s." % (str(self.curve_description))
-        return [dict(name='response', title='Direct marginal response', description=description)]
+        return [dict(name=self.label, title='Direct marginal response', description=description)]
 
     @staticmethod
     def describe():

openest/models/curve.py

@@ -36,6 +36,7 @@ def make_linear_spline_curve(xx, yy, limits):
 
 class FlatCurve(CurveCurve):
     def __init__(self, yy):
+        self.yy = yy
         super(FlatCurve, self).__init__([-np.inf, np.inf], lambda x: yy)
 
 class LinearCurve(CurveCurve):
@@ -149,7 +150,7 @@ def __call__(self, xs):
         clipping = self.clipping_curve(xs)
         ys = [y if y is not None else 0 for y in ys]
         clipping = [y if not np.isnan(y) else 0 for y in clipping]
-        return ys * (clipping > self.clipy)
+        return ys * (np.array(clipping) > self.clipy)
 
 class MinimumCurve(UnivariateCurve):
     def __init__(self, curve1, curve2):
@@ -183,6 +184,7 @@ def __call__(self, xs):
 class PiecewiseCurve(UnivariateCurve):
     def __init__(self, curves, knots, xtrans=lambda x: x):
         super(PiecewiseCurve, self).__init__(knots)
+        assert len(curves) == len(knots) - 1
         self.curves = curves
         self.knots = knots
         self.xtrans = xtrans # for example, to select first column