rtc-tools 2.5.2rc3__py3-none-any.whl → 2.6.0__py3-none-any.whl

rtctools/optimization/control_tree_mixin.py

@@ -41,14 +41,35 @@ class ControlTreeMixin(OptimizationProblem):
 
         options = {}
 
-        options['forecast_variables'] = [var.name()
-                                         for var in self.dae_variables['constant_inputs']]
-        options['branching_times'] = self.times()[1:]
-        options['k'] = 2
+        options["forecast_variables"] = [
+            var.name() for var in self.dae_variables["constant_inputs"]
+        ]
+        options["branching_times"] = self.times()[1:]
+        options["k"] = 2
 
         return options
 
+    def discretize_control(self, variable, ensemble_member, times, offset):
+        control_indices = np.zeros(len(times), dtype=np.int16)
+        for branch, members in self.__branches.items():
+            if ensemble_member not in members:
+                continue
+
+            branching_time_0 = self.__branching_times[len(branch) + 0]
+            branching_time_1 = self.__branching_times[len(branch) + 1]
+            els = np.logical_and(times >= branching_time_0, times < branching_time_1)
+            nnz = np.count_nonzero(els)
+            try:
+                control_indices[els] = self.__discretize_controls_cache[(variable, branch)]
+            except KeyError:
+                control_indices[els] = list(range(offset, offset + nnz))
+                self.__discretize_controls_cache[(variable, branch)] = control_indices[els]
+                offset += nnz
+        return control_indices
+
     def discretize_controls(self, resolved_bounds):
+        self.__discretize_controls_cache = {}
+
         # Collect options
         options = self.control_tree_options()
 
@@ -56,14 +77,14 @@ class ControlTreeMixin(OptimizationProblem):
         # presence of these is assumed below.
         times = self.times()
         t0 = self.initial_time
-        branching_times = options['branching_times']
-        n_branching_times = len(branching_times)
+        self.__branching_times = options["branching_times"]
+        n_branching_times = len(self.__branching_times)
         if n_branching_times > len(times) - 1:
             raise Exception("Too many branching points specified")
-        branching_times = np.concatenate(([t0], branching_times, [np.inf]))
+        self.__branching_times = np.concatenate(([t0], self.__branching_times, [np.inf]))
 
         logger.debug("ControlTreeMixin: Branching times:")
-        logger.debug(branching_times)
+        logger.debug(self.__branching_times)
 
         # Branches start at branching times, so that the tree looks like the following:
         #
@@ -88,8 +109,8 @@ class ControlTreeMixin(OptimizationProblem):
             distances = np.zeros((n_branch_members, n_branch_members))
 
             # Decide branching on a segment of the time horizon
-            branching_time_0 = branching_times[len(current_branch) + 1]
-            branching_time_1 = branching_times[len(current_branch) + 2]
+            branching_time_0 = self.__branching_times[len(current_branch) + 1]
+            branching_time_1 = self.__branching_times[len(current_branch) + 2]
 
             # Compute reverse ensemble member index-to-distance index map.
             reverse = {}
@@ -98,23 +119,24 @@ class ControlTreeMixin(OptimizationProblem):
 
             # Compute distances between ensemble members, summed for all
             # forecast variables
-            for forecast_variable in options['forecast_variables']:
+            for forecast_variable in options["forecast_variables"]:
                 # We assume the time stamps of the forecasts in all ensemble
                 # members to be identical
-                timeseries = self.constant_inputs(ensemble_member=0)[
-                    forecast_variable]
+                timeseries = self.constant_inputs(ensemble_member=0)[forecast_variable]
                 els = np.logical_and(
-                    timeseries.times >= branching_time_0, timeseries.times < branching_time_1)
+                    timeseries.times >= branching_time_0, timeseries.times < branching_time_1
+                )
 
                 # Compute distance between ensemble members
                 for i, member_i in enumerate(branches[current_branch]):
-                    timeseries_i = self.constant_inputs(ensemble_member=member_i)[
-                        forecast_variable]
+                    timeseries_i = self.constant_inputs(ensemble_member=member_i)[forecast_variable]
                     for j, member_j in enumerate(branches[current_branch]):
                         timeseries_j = self.constant_inputs(ensemble_member=member_j)[
-                            forecast_variable]
-                        distances[
-                            i, j] += np.linalg.norm(timeseries_i.values[els] - timeseries_j.values[els])
+                            forecast_variable
+                        ]
+                        distances[i, j] += np.linalg.norm(
+                            timeseries_i.values[els] - timeseries_j.values[els]
+                        )
 
             # Keep track of ensemble members that have not yet been allocated
             # to a new branch
@@ -123,45 +145,51 @@ class ControlTreeMixin(OptimizationProblem):
             # We first select the scenario with the max distance to any other branch
             idx = np.argmax(np.amax(distances, axis=0))
 
-            for i in range(options['k']):
+            for i in range(options["k"]):
                 if idx >= 0:
-                    branches[current_branch +
-                             (i, )] = [branches[current_branch][idx]]
+                    branches[current_branch + (i,)] = [branches[current_branch][idx]]
 
                     available.remove(branches[current_branch][idx])
 
                     # We select the scenario with the max min distance to the other branches
-                    min_distances = np.array([
-                        min([np.inf] + [distances[j, k]
-                            for j, member_j in enumerate(branches[current_branch])
-                            if member_j not in available and member_k in available])
-                        for k, member_k in enumerate(branches[current_branch])
-                    ], dtype=np.float64)
+                    min_distances = np.array(
+                        [
+                            min(
+                                [np.inf]
+                                + [
+                                    distances[j, k]
+                                    for j, member_j in enumerate(branches[current_branch])
+                                    if member_j not in available and member_k in available
+                                ]
+                            )
+                            for k, member_k in enumerate(branches[current_branch])
+                        ],
+                        dtype=np.float64,
+                    )
                     min_distances[np.where(min_distances == np.inf)] = -np.inf
 
                     idx = np.argmax(min_distances)
                     if min_distances[idx] <= 0:
                         idx = -1
                 else:
-                    branches[current_branch + (i, )] = []
+                    branches[current_branch + (i,)] = []
 
             # Cluster remaining ensemble members to branches
             for member_i in available:
                 min_i = 0
                 min_distance = np.inf
-                for i in range(options['k']):
-                    branch2 = branches[current_branch + (i, )]
+                for i in range(options["k"]):
+                    branch2 = branches[current_branch + (i,)]
                     if len(branch2) > 0:
-                        distance = distances[
-                            reverse[member_i], reverse[branch2[0]]]
+                        distance = distances[reverse[member_i], reverse[branch2[0]]]
                         if distance < min_distance:
                             min_distance = distance
                             min_i = i
-                branches[current_branch + (min_i, )].append(member_i)
+                branches[current_branch + (min_i,)].append(member_i)
 
             # Recurse
-            for i in range(options['k']):
-                branch(current_branch + (i, ))
+            for i in range(options["k"]):
+                branch(current_branch + (i,))
 
         current_branch = ()
         branches[current_branch] = list(range(self.ensemble_size))
@@ -172,36 +200,10 @@ class ControlTreeMixin(OptimizationProblem):
 
         self.__branches = branches
 
-        # Map ensemble members to control inputs
-        # (variable, (ensemble member, step)) -> control_index
-        self.__control_indices = [{} for ensemble_member in range(self.ensemble_size)]
-        count = 0
-        for control_input in self.controls:
-            times = self.times(control_input)
-            for member in range(self.ensemble_size):
-                self.__control_indices[member][control_input] = np.zeros(
-                    len(times), dtype=np.int16)
-            for branch, members in branches.items():
-                if not members:
-                    # Avoid making free variables by skipping branches which have no members
-                    continue
-
-                branching_time_0 = branching_times[len(branch) + 0]
-                branching_time_1 = branching_times[len(branch) + 1]
-                els = np.logical_and(
-                    times >= branching_time_0, times < branching_time_1)
-                nnz = np.count_nonzero(els)
-                for member in members:
-                    self.__control_indices[member][control_input][els] = \
-                        list(range(count, count + nnz))
-                count += nnz
-
-        discrete = self._collint_get_discrete(count, self.__control_indices)
-        lbx, ubx = self._collint_get_lbx_ubx(count, self.__control_indices)
-        x0 = self._collint_get_x0(count, self.__control_indices)
-
-        # Return number of control variables
-        return count, discrete, lbx, ubx, x0, self.__control_indices
+        # By now, the tree branches have been set up.  We now rely
+        # on the default discretization logic to call discretize_control()
+        # for each (control variable, ensemble member) pair.
+        return super().discretize_controls(resolved_bounds)
 
     @property
     def control_tree_branches(self) -> Dict[Tuple[int], List[int]]:

rtctools/optimization/csv_lookup_table_mixin.py

@@ -6,8 +6,9 @@ import pickle
 from typing import Iterable, List, Tuple, Union
 
 import casadi as ca
-
 import numpy as np
+from scipy.interpolate import bisplev, bisplrep, splev
+from scipy.optimize import brentq
 
 import rtctools.data.csv as csv
 from rtctools._internal.caching import cached
@@ -15,9 +16,6 @@ from rtctools.data.interpolation.bspline1d import BSpline1D
 from rtctools.data.interpolation.bspline2d import BSpline2D
 from rtctools.optimization.timeseries import Timeseries
 
-from scipy.interpolate import bisplev, bisplrep, splev
-from scipy.optimize import brentq
-
 from .optimization_problem import LookupTable as LookupTableBase
 from .optimization_problem import OptimizationProblem
 
@@ -54,10 +52,14 @@ class LookupTable(LookupTableBase):
     def domain(self) -> Tuple:
         t = self.__t
         if t is None:
-            raise AttributeError('This lookup table was not instantiated with tck metadata. \
-                                  Domain/Range information is unavailable.')
+            raise AttributeError(
+                "This lookup table was not instantiated with tck metadata. \
+                                  Domain/Range information is unavailable."
+            )
         if type(t) == tuple and len(t) == 2:
-            raise NotImplementedError('Domain/Range information is not yet implemented for 2D LookupTables')
+            raise NotImplementedError(
+                "Domain/Range information is not yet implemented for 2D LookupTables"
+            )
 
         return np.nextafter(t[0], np.inf), np.nextafter(t[-1], -np.inf)
 
@@ -84,9 +86,7 @@ class LookupTable(LookupTableBase):
     @cached
     def __numeric_function_evaluator(self):
         return np.vectorize(
-            lambda *args: np.nan
-            if np.any(np.isnan(args))
-            else float(self.function(*args))
+            lambda *args: np.nan if np.any(np.isnan(args)) else float(self.function(*args))
         )
 
     def __call__(
@@ -170,9 +170,7 @@ class LookupTable(LookupTableBase):
             all_viol = y_array_not_nan[lb_viol | ub_viol]
             if all_viol.size > 0:
                 raise ValueError(
-                    "Values {} are not in lookup table range ({}, {})".format(
-                        all_viol, l_r, u_r
-                    )
+                    "Values {} are not in lookup table range ({}, {})".format(all_viol, l_r, u_r)
                 )
 
         # Construct function to do inverse evaluation
@@ -200,9 +198,9 @@ class CSVLookupTableMixin(OptimizationProblem):
     """
     Adds lookup tables to your optimization problem.
 
-    During preprocessing, the CSV files located inside the ``lookup_tables`` subfolder are read.
-    In every CSV file, the first column contains the output of the lookup table.  Subsequent columns contain
-    the input variables.
+    During preprocessing, the CSV files located inside the ``lookup_tables`` subfolder are read. In
+    every CSV file, the first column contains the output of the lookup table. Subsequent columns
+    contain the input variables.
 
     Cubic B-Splines are used to turn the data points into continuous lookup tables.
 
@@ -226,13 +224,15 @@ class CSVLookupTableMixin(OptimizationProblem):
         Currently only one-dimensional lookup tables are fully supported.  Support for two-
         dimensional lookup tables is experimental.
 
-    :cvar csv_delimiter:                 Column delimiter used in CSV files.  Default is ``,``.
-    :cvar csv_lookup_table_debug:        Whether to generate plots of the spline fits.  Default is ``false``.
-    :cvar csv_lookup_table_debug_points: Number of evaluation points for plots.  Default is ``100``.
+    :cvar csv_delimiter:  Column delimiter used in CSV files. Default is ``,``.
+    :cvar csv_lookup_table_debug:  Whether to generate plots of the spline fits.
+        Default is ``False``.
+    :cvar csv_lookup_table_debug_points:  Number of evaluation points for plots.
+        Default is ``100``.
     """
 
     #: Column delimiter used in CSV files
-    csv_delimiter = ','
+    csv_delimiter = ","
 
     #: Debug settings
     csv_lookup_table_debug = False
@@ -240,13 +240,12 @@ class CSVLookupTableMixin(OptimizationProblem):
 
     def __init__(self, **kwargs):
         # Check arguments
-        if 'input_folder' in kwargs:
-            assert ('lookup_table_folder' not in kwargs)
+        if "input_folder" in kwargs:
+            assert "lookup_table_folder" not in kwargs
 
-            self.__lookup_table_folder = os.path.join(
-                kwargs['input_folder'], 'lookup_tables')
+            self.__lookup_table_folder = os.path.join(kwargs["input_folder"], "lookup_tables")
         else:
-            self.__lookup_table_folder = kwargs['lookup_table_folder']
+            self.__lookup_table_folder = kwargs["lookup_table_folder"]
 
         # Call parent
         super().__init__(**kwargs)
@@ -256,15 +255,15 @@ class CSVLookupTableMixin(OptimizationProblem):
         super().pre()
 
         # Get curve fitting options from curvefit_options.ini file
-        ini_path = os.path.join(
-            self.__lookup_table_folder, 'curvefit_options.ini')
+        ini_path = os.path.join(self.__lookup_table_folder, "curvefit_options.ini")
         try:
             ini_config = configparser.RawConfigParser()
             ini_config.read(ini_path)
             no_curvefit_options = False
         except IOError:
             logger.info(
-                "CSVLookupTableMixin: No curvefit_options.ini file found. Using default values.")
+                "CSVLookupTableMixin: No curvefit_options.ini file found. Using default values."
+            )
             no_curvefit_options = True
 
         def get_curvefit_options(curve_name, no_curvefit_options=no_curvefit_options):
@@ -282,30 +281,33 @@ class CSVLookupTableMixin(OptimizationProblem):
                     prop = 0
                 except ValueError:
                     raise Exception(
-                        'CSVLookupTableMixin: Invalid {0} constraint for {1}. {0} should '
-                        'be either -1, 0, or 1.'.format(prop_name, curve_name))
+                        "CSVLookupTableMixin: Invalid {0} constraint for {1}. {0} should "
+                        "be either -1, 0, or 1.".format(prop_name, curve_name)
+                    )
                 return prop
 
-            for prop_name in ['monotonicity', 'monotonicity2', 'curvature']:
+            for prop_name in ["monotonicity", "monotonicity2", "curvature"]:
                 curvefit_options.append(get_property(prop_name))
 
-            logger.debug("CSVLookupTableMixin: Curve fit option for {}:({},{},{})".format(
-                curve_name, *curvefit_options))
+            logger.debug(
+                "CSVLookupTableMixin: Curve fit option for {}:({},{},{})".format(
+                    curve_name, *curvefit_options
+                )
+            )
             return tuple(curvefit_options)
 
         def check_lookup_table(lookup_table):
             if lookup_table in self.__lookup_tables:
                 raise Exception(
                     "Cannot add lookup table {},"
-                    "since there is already one with this name.".format(lookup_table))
+                    "since there is already one with this name.".format(lookup_table)
+                )
+
         # Read CSV files
-        logger.info(
-            "CSVLookupTableMixin: Generating Splines from lookup table data.")
+        logger.info("CSVLookupTableMixin: Generating Splines from lookup table data.")
         self.__lookup_tables = {}
         for filename in glob.glob(os.path.join(self.__lookup_table_folder, "*.csv")):
-
-            logger.debug(
-                "CSVLookupTableMixin: Reading lookup table from {}".format(filename))
+            logger.debug("CSVLookupTableMixin: Reading lookup table from {}".format(filename))
 
             csvinput = csv.load(filename, delimiter=self.csv_delimiter)
             output = csvinput.dtype.names[0]
@@ -314,62 +316,80 @@ class CSVLookupTableMixin(OptimizationProblem):
             # Get monotonicity and curvature from ini file
             mono, mono2, curv = get_curvefit_options(output)
 
-            logger.debug(
-                "CSVLookupTableMixin: Output is {}, inputs are {}.".format(output, inputs))
+            logger.debug("CSVLookupTableMixin: Output is {}, inputs are {}.".format(output, inputs))
 
             tck = None
             function = None
 
-            # If tck file is newer than the csv file, first try to load the cached values from the tck file
-            tck_filename = filename.replace('.csv', '.tck')
+            # If tck file is newer than the csv file, first try to load the cached values from
+            # the tck file
+            tck_filename = filename.replace(".csv", ".tck")
             valid_cache = False
             if os.path.exists(tck_filename):
                 if no_curvefit_options:
                     valid_cache = os.path.getmtime(filename) < os.path.getmtime(tck_filename)
                 else:
-                    valid_cache = (os.path.getmtime(filename) < os.path.getmtime(tck_filename)) and \
-                                  (os.path.getmtime(ini_path) < os.path.getmtime(tck_filename))
+                    valid_cache = (
+                        os.path.getmtime(filename) < os.path.getmtime(tck_filename)
+                    ) and (os.path.getmtime(ini_path) < os.path.getmtime(tck_filename))
                 if valid_cache:
                     logger.debug(
-                        'CSVLookupTableMixin: Attempting to use cached tck values for {}'.format(output))
-                    with open(tck_filename, 'rb') as f:
+                        "CSVLookupTableMixin: Attempting to use cached tck values for {}".format(
+                            output
+                        )
+                    )
+                    with open(tck_filename, "rb") as f:
                         try:
                             tck, function = pickle.load(f)
                         except Exception:
                             valid_cache = False
             if not valid_cache:
-                logger.info(
-                    'CSVLookupTableMixin: Recalculating tck values for {}'.format(output))
+                logger.info("CSVLookupTableMixin: Recalculating tck values for {}".format(output))
 
             if len(csvinput.dtype.names) == 2:
                 if not valid_cache:
                     k = 3  # default value
                     # 1D spline fitting needs k+1 data points
                     if len(csvinput[output]) >= k + 1:
-                        tck = BSpline1D.fit(csvinput[inputs[0]], csvinput[
-                                            output], k=k, monotonicity=mono, curvature=curv)
+                        tck = BSpline1D.fit(
+                            csvinput[inputs[0]],
+                            csvinput[output],
+                            k=k,
+                            monotonicity=mono,
+                            curvature=curv,
+                        )
                     else:
                         raise Exception(
-                            'CSVLookupTableMixin: Too few data points in {} to do spline fitting. '
-                            'Need at least {} points.'.format(filename, k + 1))
+                            "CSVLookupTableMixin: Too few data points in {} to do spline fitting. "
+                            "Need at least {} points.".format(filename, k + 1)
+                        )
 
                 if self.csv_lookup_table_debug:
                     import pylab
-                    i = np.linspace(csvinput[inputs[0]][0], csvinput[
-                                    inputs[0]][-1], self.csv_lookup_table_debug_points)
+
+                    i = np.linspace(
+                        csvinput[inputs[0]][0],
+                        csvinput[inputs[0]][-1],
+                        self.csv_lookup_table_debug_points,
+                    )
                     o = splev(i, tck)
                     pylab.clf()
                     # TODO: Figure out why cross-section B0607 in NZV does not
                     # conform to constraints!
                     pylab.plot(i, o)
-                    pylab.plot(csvinput[inputs[0]], csvinput[
-                               output], linestyle='', marker='x', markersize=10)
-                    figure_filename = filename.replace('.csv', '.png')
+                    pylab.plot(
+                        csvinput[inputs[0]],
+                        csvinput[output],
+                        linestyle="",
+                        marker="x",
+                        markersize=10,
+                    )
+                    figure_filename = filename.replace(".csv", ".png")
                     pylab.savefig(figure_filename)
 
                 symbols = [ca.SX.sym(inputs[0])]
                 if not valid_cache:
-                    function = ca.Function('f', symbols, [BSpline1D(*tck)(symbols[0])])
+                    function = ca.Function("f", symbols, [BSpline1D(*tck)(symbols[0])])
                 check_lookup_table(output)
                 self.__lookup_tables[output] = LookupTable(symbols, function, tck)
 
@@ -382,42 +402,55 @@ class CSVLookupTableMixin(OptimizationProblem):
                     if len(csvinput[output]) >= (kx + 1) * (ky + 1):
                         # TODO: add curvature parameters from curvefit_options.ini
                         # once 2d fit is implemented
-                        tck = bisplrep(csvinput[inputs[0]], csvinput[
-                                       inputs[1]], csvinput[output], kx=kx, ky=ky)
+                        tck = bisplrep(
+                            csvinput[inputs[0]], csvinput[inputs[1]], csvinput[output], kx=kx, ky=ky
+                        )
                     else:
                         raise Exception(
-                            'CSVLookupTableMixin: Too few data points in {} to do spline fitting. '
-                            'Need at least {} points.'.format(filename, (kx + 1) * (ky + 1)))
+                            "CSVLookupTableMixin: Too few data points in {} to do spline fitting. "
+                            "Need at least {} points.".format(filename, (kx + 1) * (ky + 1))
+                        )
 
                 if self.csv_lookup_table_debug:
                     import pylab
-                    i1 = np.linspace(csvinput[inputs[0]][0], csvinput[
-                                     inputs[0]][-1], self.csv_lookup_table_debug_points)
-                    i2 = np.linspace(csvinput[inputs[1]][0], csvinput[
-                                     inputs[1]][-1], self.csv_lookup_table_debug_points)
+
+                    i1 = np.linspace(
+                        csvinput[inputs[0]][0],
+                        csvinput[inputs[0]][-1],
+                        self.csv_lookup_table_debug_points,
+                    )
+                    i2 = np.linspace(
+                        csvinput[inputs[1]][0],
+                        csvinput[inputs[1]][-1],
+                        self.csv_lookup_table_debug_points,
+                    )
                     i1, i2 = np.meshgrid(i1, i2)
                     i1 = i1.flatten()
                     i2 = i2.flatten()
                     o = bisplev(i1, i2, tck)
                     pylab.clf()
                     pylab.plot_surface(i1, i2, o)
-                    figure_filename = filename.replace('.csv', '.png')
+                    figure_filename = filename.replace(".csv", ".png")
                     pylab.savefig(figure_filename)
                 symbols = [ca.SX.sym(inputs[0]), ca.SX.sym(inputs[1])]
                 if not valid_cache:
-                    function = ca.Function('f', symbols, [BSpline2D(*tck)(symbols[0], symbols[1])])
+                    function = ca.Function("f", symbols, [BSpline2D(*tck)(symbols[0], symbols[1])])
                 check_lookup_table(output)
                 self.__lookup_tables[output] = LookupTable(symbols, function, tck)
 
             else:
                 raise Exception(
-                    'CSVLookupTableMixin: {}-dimensional lookup tables not implemented yet.'.format(
-                        len(csvinput.dtype.names)))
+                    "CSVLookupTableMixin: {}-dimensional lookup tables not implemented yet.".format(
+                        len(csvinput.dtype.names)
+                    )
+                )
 
             if not valid_cache:
-                pickle.dump((tck, function),
-                            open(filename.replace('.csv', '.tck'), 'wb'),
-                            protocol=pickle.HIGHEST_PROTOCOL)
+                pickle.dump(
+                    (tck, function),
+                    open(filename.replace(".csv", ".tck"), "wb"),
+                    protocol=pickle.HIGHEST_PROTOCOL,
+                )
 
     def lookup_tables(self, ensemble_member):
         # Call parent class first for default values.