Add data sampling class

docs/source/custom_envs.rst

@@ -1,5 +1,5 @@
 Create Custom Environments
-===========================
+==========================
 
 By inheriting from the :code:`OpfEnv` base class, a wide variety of custom 
 environments can be created. In the process, some steps have to be considered.

docs/source/index.rst

@@ -42,6 +42,7 @@ The repository can be found on
    benchmarks
    api_base_class
    environment_design
+   sampling 
    custom_envs
    advanced_features
    supervised_learning

docs/source/sampling.rst

@@ -0,0 +1,49 @@
+Data Sampling
+=============
+
+The DataSampler Class
+---------------------
+
+In the gymnasium :meth:`reset` method, the environment is set to some random 
+state. This is done by the :meth:`_sampling` method of the base class, which
+calls a :class:`DataSampler` object. In *OPF-Gym*, three standard data samplers
+are pre-implemented: The :class:`SimbenchSampler`, the :class:`NormalSampler`,
+and the :class:`UniformSampler`. The default in *OPF-Gym* is to use SimBench 
+data for active and reactive power values, and a uniform distribution for state
+variables that are not included in the SimBench data (e.g. prices, slack 
+voltage, etc.).
+
+SimBench Data Sampler 
+_____________________
+.. autoclass:: opfgym.sampling.SimbenchSampler
+   :members:
+
+Normal Distribution Data Sampler
+_______________________________
+.. autoclass:: opfgym.sampling.NormalSampler
+   :members:
+
+Uniform Distribution Data Sampler
+_______________________________
+.. autoclass:: opfgym.sampling.UniformSampler
+   :members:
+
+
+Data Sampling Wrappers
+----------------------
+
+In many cases, we want to combine multiple distributions for different state
+variables, for example by sampling generation data from one distribution
+and market prices from another. In *OPF-Gym*, this is done with the 
+:class:`DataSamplerWrapper` class. Two standard wrappers are pre-implemented:
+The :class:`SequentialSampler` and the :class:`MixedRandomSampler`.
+
+The SequentialSampler
+_______________________________
+.. autoclass:: opfgym.sampling.SequentialSampler
+   :members:
+
+The MixedRandomSampler
+_______________________________
+.. autoclass:: opfgym.sampling.MixedRandomSampler
+   :members:

opfgym/envs/eco_dispatch.py

@@ -24,7 +24,7 @@ class EcoDispatch(opf_env.OpfEnv):
 
     """
 
-    def __init__(self, simbench_network_name='1-HV-urban--0-sw', 
+    def __init__(self, simbench_network_name='1-HV-urban--0-sw',
                  gen_scaling=1.0, load_scaling=1.5, max_price_eur_gwh=0.5,
                  min_power=0, *args, **kwargs):
 
@@ -41,14 +41,16 @@ def __init__(self, simbench_network_name='1-HV-urban--0-sw',
 
         # Define the RL problem
         # See all load power values, non-controlled generators, and generator prices...
-        obs_keys = [('load', 'p_mw', net.load.index),
-                    ('load', 'q_mvar', net.load.index),
-                    ('poly_cost', 'cp1_eur_per_mw', net.poly_cost.index),
-                    ('pwl_cost', 'cp1_eur_per_mw', net.pwl_cost.index),
-                    # These 3 are not relevant because len=0, if the default is used
-                    ('sgen', 'p_mw', net.sgen.index[~net.sgen.controllable]),
-                    ('storage', 'p_mw', net.storage.index),
-                    ('storage', 'q_mvar', net.storage.index)]
+        obs_keys = [
+            ('load', 'p_mw', net.load.index),
+            ('load', 'q_mvar', net.load.index),
+            ('poly_cost', 'cp1_eur_per_mw', net.poly_cost.index),
+            ('pwl_cost', 'cp1_eur_per_mw', net.pwl_cost.index),
+            # These 3 are not relevant because len=0, if the default is used
+            ('sgen', 'p_mw', net.sgen.index[~net.sgen.controllable]),
+            ('storage', 'p_mw', net.storage.index),
+            ('storage', 'q_mvar', net.storage.index)
+        ]
 
         # ... and control all generators' active power values
         act_keys = [('sgen', 'p_mw', net.sgen.index[net.sgen.controllable]),
@@ -89,14 +91,14 @@ def _define_opf(self, simbench_network_name, *args, **kwargs):
 
         # Add price params to the network (as poly cost so that the OPF works)
         # Note that the external grids are seen as normal power plants
-        for idx in net.ext_grid.index:
-            # Use piece-wise linear costs to prevent negative costs for negative
-            # power, which would incentivize a constraint violation (see above)
-            pp.create_pwl_cost(net, idx, 'ext_grid', points=[[0, 10000, 1]])
         for idx in net.sgen.index[net.sgen.controllable]:
             pp.create_poly_cost(net, idx, 'sgen', cp1_eur_per_mw=0)
         for idx in net.gen.index[net.gen.controllable]:
             pp.create_poly_cost(net, idx, 'gen', cp1_eur_per_mw=0)
+        # Use piece-wise linear costs to prevent negative costs for negative
+        # power, which would incentivize a constraint violation (see above)
+        for idx in net.ext_grid.index:
+            pp.create_pwl_cost(net, idx, 'ext_grid', points=[[0, 10000, 1]])
 
         net.poly_cost['min_cp1_eur_per_mw'] = 0
         net.poly_cost['max_cp1_eur_per_mw'] = self.max_price_eur_gwh
@@ -111,13 +113,8 @@ def _define_opf(self, simbench_network_name, *args, **kwargs):
     def _sampling(self, *args, **kwargs):
         super()._sampling(*args, **kwargs)
 
-        # Sample prices uniformly from min/max range for gens/sgens/ext_grids
-        self._sample_from_range(
-            'poly_cost', 'cp1_eur_per_mw', self.net.poly_cost.index)
-        self._sample_from_range(
-            'pwl_cost', 'cp1_eur_per_mw', self.net.pwl_cost.index)
-
-        # Manually update the costs in the pwl 'points' definition
+        # Manually update the costs in the pwl 'points' definition so that it's
+        # usable for the pandapower OPF
         for idx in self.net.ext_grid.index:
             price = self.net.pwl_cost.at[idx, 'cp1_eur_per_mw']
             self.net.pwl_cost.at[idx, 'points'] = [[0, 10000, price]]

opfgym/envs/load_shedding.py

@@ -60,8 +60,8 @@ def __init__(self, simbench_network_name='1-MV-comm--2-sw',
             ('load', 'p_mw', net.load.index),
             ('load', 'q_mvar', net.load.index),
             ('storage', 'p_mw', net.storage.index[~net.storage.controllable]),
-            # ('poly_cost', 'cp1_eur_per_mw', net.poly_cost.index),  # Separately sampled in _sampling(), see below
-            # ('pwl_cost', 'cp1_eur_per_mw', net.pwl_cost.index)
+            ('poly_cost', 'cp1_eur_per_mw', net.poly_cost.index),
+            ('pwl_cost', 'cp1_eur_per_mw', net.pwl_cost.index)
         ]
 
         # Control active power of loads and storages
@@ -122,14 +122,6 @@ def _define_opf(self, simbench_network_name, *args, **kwargs):
     def _sampling(self, *args, **kwargs):
         super()._sampling(*args, **kwargs)
 
-        # Sample prices for loads and storages
-        # The idea is that not always the same loads should be shedded. Instead,
-        # the current situation should be considered, represented by some price.
-        self._sample_from_range(
-            'poly_cost', 'cp1_eur_per_mw', self.net.poly_cost.index)
-        self._sample_from_range(
-            'pwl_cost', 'cp1_eur_per_mw', self.net.pwl_cost.index)
-
         # Manually update the points of the piece-wise linear costs for storage
         for idx in self.net.pwl_cost.index:
             price = self.net.pwl_cost.at[idx, 'cp1_eur_per_mw']

opfgym/envs/voltage_control.py

@@ -111,13 +111,6 @@ def _define_opf(self, simbench_network_name, *args, **kwargs):
     def _sampling(self, *args, **kwargs):
         super()._sampling(*args, **kwargs)
 
-        # Sample reactive power prices uniformly from min/max range
-        if self.market_based:
-            for unit_type in ('sgen', 'ext_grid', 'storage'):
-                self._sample_from_range(
-                'poly_cost', 'cq2_eur_per_mvar2',
-                self.net.poly_cost[self.net.poly_cost.et == unit_type].index)
-
         # Active power is not controllable (only relevant for OPF baseline)
         # Set active power boundaries to current active power values
         for unit_type in ('sgen', 'storage'):

opfgym/examples/mixed_continuous_discrete.py

@@ -74,10 +74,6 @@ def _define_opf(self, simbench_network_name, *args, **kwargs):
     def _sampling(self, *args, **kwargs):
         super()._sampling(*args, **kwargs)
 
-        # Sample slack voltage randomly to make the problem more difficult
-        # so that trafo tap changing is required for voltage control
-        self._sample_from_range('ext_grid', 'vm_pu', self.net.ext_grid.index)
-
         # Active power is not controllable (only relevant for OPF baseline)
         # Set active power boundaries to current active power values
         for unit_type in ('sgen',):

opfgym/examples/multi_stage.py

@@ -19,12 +19,14 @@
 
 class MultiStageOpf(MultiStageOpfEnv):
     def __init__(self, simbench_network_name='1-LV-urban6--0-sw',
-                 steps_per_episode=4, train_data='simbench',
-                 test_data='simbench',
+                 steps_per_episode=4, train_sampling='simbench',
+                 test_sampling='simbench', validation_sampling='simbench',
                  *args, **kwargs):
 
         assert steps_per_episode > 1, "At least two steps required for a multi-stage OPF."
-        assert 'simbench' in train_data and 'simbench' in test_data, "Only simbench networks are supported because time-series data required."
+        assert 'simbench' in train_sampling, "Only simbench networks are supported because time-series data required."
+        assert 'simbench' in test_sampling, "Only simbench networks are supported because time-series data required."
+        assert 'simbench' in validation_sampling, "Only simbench networks are supported because time-series data required."
 
         net, profiles = self._define_opf(
             simbench_network_name, *args, **kwargs)
@@ -41,6 +43,9 @@ def __init__(self, simbench_network_name='1-LV-urban6--0-sw',
         super().__init__(net, act_keys, obs_keys, profiles=profiles,
                          steps_per_episode=steps_per_episode,
                          optimal_power_flow_solver=False,
+                         train_sampling=train_sampling,
+                         test_sampling=test_sampling,
+                         validation_sampling=validation_sampling,
                          *args, **kwargs)
 
     def _define_opf(self, simbench_network_name, *args, **kwargs):

opfgym/examples/non_simbench_net.py

@@ -11,11 +11,15 @@
 
 
 class NonSimbenchNet(opf_env.OpfEnv):
-    def __init__(self, train_data='normal_around_mean',
-                 test_data='normal_around_mean',
+    def __init__(self,
+                 train_sampling='normal_around_mean',
+                 test_sampling='normal_around_mean',
+                 validation_sampling='normal_around_mean',
                  *args, **kwargs):
 
-        assert 'simbench' not in train_data and 'simbench' not in test_data, "Only non-simbench networks are supported."
+        assert 'simbench' not in train_sampling, "Only non-simbench networks are supported."
+        assert 'simbench' not in test_sampling, "Only non-simbench networks are supported."
+        assert 'simbench' not in validation_sampling, "Only non-simbench networks are supported."
 
         net = self._define_opf()
 
@@ -30,7 +34,9 @@ def __init__(self, train_data='normal_around_mean',
         act_keys = [('gen', 'p_mw', net.gen.index)]
 
         super().__init__(net, act_keys, obs_keys,
-                         train_data=train_data, test_data=test_data,
+                         train_sampling=train_sampling,
+                         test_sampling=test_sampling,
+                         validation_sampling=validation_sampling,
                          *args, **kwargs)
 
     def _define_opf(self):

opfgym/multi_stage.py

@@ -16,18 +16,27 @@ class MultiStageOpfEnv(OpfEnv):
         Same as the base class
 
     """
-    def __init__(self, *args, steps_per_episode: int=4, **kwargs):
+    def __init__(self, *args, steps_per_episode: int = 4, **kwargs):
         assert steps_per_episode > 1, "At least two steps required for a multi-stage OPF."
-        if kwargs.get('train_data') and isinstance(kwargs.get('train_data')):
-            assert 'simbench' in kwargs.get('train_data')
+        assert 'simbench' in kwargs.get('train_sampling', 'simbench')
+        assert 'simbench' in kwargs.get('validation_sampling', 'simbench')
+        assert 'simbench' in kwargs.get('test_sampling', 'simbench')
         super().__init__(*args, steps_per_episode=steps_per_episode, **kwargs)
 
 
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+        if options:
+            self.test = options.get('test', False)
+        else:
+            self.test = False
+
+
     def step(self, action):
         """ Extend step method to sample the next time step of the simbench data. """
         obs, reward, terminated, truncated, info = super().step(action)
 
-        new_step = self.current_simbench_step + 1
+        new_step = self.current_time_step + 1
 
         # Enforce train/test-split
         if self.test:
@@ -47,7 +56,7 @@ def step(self, action):
             return obs, reward, terminated, truncated, info
 
         # Increment the time-series states
-        self._sampling(step=new_step)
+        self._sampling(step=new_step, test=self.test)
 
         # Rerun the power flow calculation for the new state if required
         if self.pf_for_obs is True: