foxes 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of foxes might be problematic. Click here for more details.

Files changed (65) hide show
  1. examples/field_data_nc/run.py +11 -4
  2. examples/streamline_wakes/run.py +6 -3
  3. foxes/algorithms/downwind/downwind.py +1 -0
  4. foxes/config/__init__.py +1 -1
  5. foxes/config/config.py +87 -14
  6. foxes/constants.py +12 -1
  7. foxes/core/algorithm.py +13 -8
  8. foxes/core/engine.py +30 -0
  9. foxes/core/farm_controller.py +41 -24
  10. foxes/core/states.py +1 -1
  11. foxes/core/wind_farm.py +109 -0
  12. foxes/engines/dask.py +88 -4
  13. foxes/engines/default.py +45 -2
  14. foxes/engines/mpi.py +5 -16
  15. foxes/engines/multiprocess.py +1 -10
  16. foxes/engines/numpy.py +30 -0
  17. foxes/engines/pool.py +48 -0
  18. foxes/engines/ray.py +1 -1
  19. foxes/engines/single.py +30 -0
  20. foxes/input/farm_layout/from_csv.py +2 -2
  21. foxes/input/farm_layout/from_file.py +2 -2
  22. foxes/input/farm_layout/from_json.py +2 -2
  23. foxes/input/states/__init__.py +0 -1
  24. foxes/input/states/create/random_abl_states.py +2 -2
  25. foxes/input/states/field_data_nc.py +286 -141
  26. foxes/input/states/multi_height.py +3 -3
  27. foxes/input/states/states_table.py +3 -3
  28. foxes/input/yaml/dict.py +83 -46
  29. foxes/input/yaml/windio/__init__.py +2 -1
  30. foxes/input/yaml/windio/read_attributes.py +17 -34
  31. foxes/input/yaml/windio/read_farm.py +57 -3
  32. foxes/input/yaml/windio/read_outputs.py +116 -56
  33. foxes/input/yaml/windio/{get_states.py → read_site.py} +69 -0
  34. foxes/input/yaml/windio/windio.py +42 -119
  35. foxes/input/yaml/yaml.py +3 -3
  36. foxes/models/model_book.py +1 -0
  37. foxes/models/point_models/__init__.py +1 -0
  38. foxes/models/point_models/ustar2ti.py +84 -0
  39. foxes/models/turbine_models/lookup_table.py +2 -2
  40. foxes/models/turbine_models/sector_management.py +2 -2
  41. foxes/models/turbine_models/table_factors.py +2 -2
  42. foxes/models/turbine_types/CpCt_file.py +2 -2
  43. foxes/models/turbine_types/CpCt_from_two.py +3 -3
  44. foxes/models/turbine_types/PCt_file.py +2 -2
  45. foxes/models/turbine_types/PCt_from_two.py +3 -3
  46. foxes/models/turbine_types/TBL_file.py +2 -2
  47. foxes/models/turbine_types/wsrho2PCt_from_two.py +3 -3
  48. foxes/models/turbine_types/wsti2PCt_from_two.py +3 -3
  49. foxes/output/__init__.py +1 -0
  50. foxes/output/output.py +5 -3
  51. foxes/output/slice_data.py +1 -1
  52. foxes/output/slices_data.py +323 -0
  53. foxes/output/state_turbine_table.py +11 -0
  54. foxes/utils/__init__.py +1 -0
  55. foxes/utils/load.py +12 -4
  56. foxes/utils/wrg_utils.py +79 -0
  57. foxes/utils/xarray_utils.py +14 -3
  58. foxes/variables.py +5 -0
  59. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/METADATA +6 -2
  60. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/RECORD +64 -62
  61. foxes/input/states/slice_data_nc.py +0 -687
  62. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/LICENSE +0 -0
  63. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/WHEEL +0 -0
  64. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/entry_points.txt +0 -0
  65. {foxes-1.2.2.dist-info → foxes-1.2.4.dist-info}/top_level.txt +0 -0
@@ -60,7 +60,10 @@ if __name__ == "__main__":
60
60
  type=int,
61
61
  )
62
62
  parser.add_argument(
63
- "-npl", "--no_pre_load", help="Do not pre-load data", action="store_true"
63
+ "-lm",
64
+ "--load_mode",
65
+ help="Do load mode",
66
+ default="preload",
64
67
  )
65
68
  parser.add_argument(
66
69
  "-nf", "--nofig", help="Do not show figures", action="store_true"
@@ -70,9 +73,7 @@ if __name__ == "__main__":
70
73
  states = foxes.input.states.FieldDataNC(
71
74
  args.file_pattern,
72
75
  output_vars=[FV.WS, FV.WD, FV.TI, FV.RHO],
73
- # var2ncvar={FV.WS: "ws", FV.WD: "wd", FV.TI: "ti"},
74
- fixed_vars={FV.RHO: 1.225},
75
- pre_load=not args.no_pre_load,
76
+ load_mode=args.load_mode,
76
77
  )
77
78
 
78
79
  mbook = foxes.models.ModelBook()
@@ -115,6 +116,12 @@ if __name__ == "__main__":
115
116
  fr = farm_results.to_dataframe()
116
117
  print(fr[[FV.WD, FV.AMB_REWS, FV.REWS, FV.AMB_P, FV.P]])
117
118
 
119
+ o = foxes.output.SlicesData(algo, farm_results)
120
+ ds = o.get_states_data_xy(
121
+ z_list=[90, 100], variables=[FV.WS], resolution=50, verbosity=1
122
+ )
123
+ print(ds)
124
+
118
125
  if not args.nofig:
119
126
  o = foxes.output.FlowPlots2D(algo, farm_results)
120
127
  o.get_mean_fig_xy(FV.WS, resolution=10)
@@ -43,9 +43,12 @@ if __name__ == "__main__":
43
43
  "-nt", "--n_turbines", help="The number of turbines", default=9, type=int
44
44
  )
45
45
  parser.add_argument(
46
- "-npl", "--no_pre_load", help="Pre-load the nc data", action="store_true"
46
+ "-lm",
47
+ "--load_mode",
48
+ help="Do load mode",
49
+ default="preload",
47
50
  )
48
- parser.add_argument("-e", "--engine", help="The engine", default=None)
51
+ parser.add_argument("-e", "--engine", help="The engine", default="process")
49
52
  parser.add_argument(
50
53
  "-n", "--n_cpus", help="The number of cpus", default=None, type=int
51
54
  )
@@ -82,7 +85,7 @@ if __name__ == "__main__":
82
85
  output_vars=[FV.WS, FV.WD, FV.TI, FV.RHO],
83
86
  var2ncvar={FV.WS: "ws", FV.WD: "wd"},
84
87
  fixed_vars={FV.RHO: 1.225, FV.TI: 0.1},
85
- pre_load=not args.no_pre_load,
88
+ load_mode=args.load_mode,
86
89
  bounds_error=False,
87
90
  )
88
91
 
@@ -188,6 +188,7 @@ class Downwind(Algorithm):
188
188
 
189
189
  self.__farm_controller = self.mbook.farm_controllers.get_item(farm_controller)
190
190
  self.farm_controller.name = farm_controller
191
+ self.farm_controller.find_turbine_types(self)
191
192
 
192
193
  @property
193
194
  def states(self):
foxes/config/__init__.py CHANGED
@@ -1 +1 @@
1
- from .config import Config, config, get_path, get_output_path
1
+ from .config import Config, config, get_path, get_input_path, get_output_path
foxes/config/config.py CHANGED
@@ -1,7 +1,8 @@
1
1
  import numpy as np
2
2
  from pathlib import Path
3
+ from sys import version_info
3
4
 
4
- from foxes.utils.dict import Dict
5
+ from foxes.utils import Dict, import_module
5
6
  import foxes.constants as FC
6
7
 
7
8
 
@@ -19,11 +20,17 @@ class Config(Dict):
19
20
  FC.DTYPE: np.float64,
20
21
  FC.ITYPE: np.int64,
21
22
  FC.WORK_DIR: Path("."),
22
- FC.OUT_DIR: Path("."),
23
+ FC.INPUT_DIR: None,
24
+ FC.OUTPUT_DIR: None,
25
+ FC.NC_ENGINE: "h5netcdf",
23
26
  },
24
27
  name="config",
25
28
  )
26
29
 
30
+ # special treat for Python 3.8:
31
+ if version_info[0] == 3 and version_info[1] == 8:
32
+ self["nc_engine"] = None
33
+
27
34
  @property
28
35
  def dtype_double(self):
29
36
  """
@@ -62,12 +69,33 @@ class Config(Dict):
62
69
 
63
70
  """
64
71
  pth = self.get_item(FC.WORK_DIR)
65
- if not isinstance(pth, Path):
72
+ if self[FC.WORK_DIR] is None:
73
+ self[FC.WORK_DIR] = Path(".")
74
+ elif not isinstance(pth, Path):
66
75
  self[FC.WORK_DIR] = Path(pth)
67
76
  return self[FC.WORK_DIR]
68
77
 
69
78
  @property
70
- def out_dir(self):
79
+ def input_dir(self):
80
+ """
81
+ The input base directory
82
+
83
+ Returns
84
+ -------
85
+ pth: pathlib.Path
86
+ Path to the input base directory
87
+
88
+ """
89
+ if self[FC.INPUT_DIR] is None:
90
+ return self.work_dir
91
+ else:
92
+ pth = self.get_item(FC.INPUT_DIR)
93
+ if not isinstance(pth, Path):
94
+ self[FC.INPUT_DIR] = Path(pth)
95
+ return self[FC.INPUT_DIR]
96
+
97
+ @property
98
+ def output_dir(self):
71
99
  """
72
100
  The default output directory
73
101
 
@@ -77,7 +105,29 @@ class Config(Dict):
77
105
  Path to the default output directory
78
106
 
79
107
  """
80
- return get_path(self.get_item(FC.OUT_DIR))
108
+ if self[FC.OUTPUT_DIR] is None:
109
+ return self.work_dir
110
+ else:
111
+ pth = self.get_item(FC.OUTPUT_DIR)
112
+ if not isinstance(pth, Path):
113
+ self[FC.OUTPUT_DIR] = Path(pth)
114
+ return self[FC.OUTPUT_DIR]
115
+
116
+ @property
117
+ def nc_engine(self):
118
+ """
119
+ The NetCDF engine
120
+
121
+ Returns
122
+ -------
123
+ nce: str
124
+ The NetCDF engine
125
+
126
+ """
127
+ nce = self[FC.NC_ENGINE]
128
+ if nce is not None:
129
+ import_module(nce)
130
+ return nce
81
131
 
82
132
 
83
133
  config = Config()
@@ -86,10 +136,37 @@ config = Config()
86
136
  """
87
137
 
88
138
 
89
- def get_path(pth):
139
+ def get_path(pth, base):
140
+ """
141
+ Gets path object, respecting the base directory
142
+
143
+ Parameters
144
+ ----------
145
+ pth: str or pathlib.Path
146
+ The path, optionally relative to base
147
+ base: pathlib.Path
148
+ The base directory
149
+
150
+ Returns
151
+ -------
152
+ out: pathlib.Path
153
+ The path, absolute or relative to base directory
154
+
155
+ :group: foxes.config
156
+
157
+ """
158
+ if not isinstance(pth, Path):
159
+ pth = Path(pth)
160
+ if pth.is_absolute():
161
+ return pth
162
+ else:
163
+ return base / pth
164
+
165
+
166
+ def get_input_path(pth):
90
167
  """
91
168
  Gets path object, respecting the configurations
92
- work directory
169
+ input directory
93
170
 
94
171
  Parameters
95
172
  ----------
@@ -99,15 +176,13 @@ def get_path(pth):
99
176
  Returns
100
177
  -------
101
178
  out: pathlib.Path
102
- The path, absolute or relative to working directory
179
+ The path, absolute or relative to input directory
103
180
  from config
104
181
 
105
182
  :group: foxes.config
106
183
 
107
184
  """
108
- if not isinstance(pth, Path):
109
- pth = Path(pth)
110
- return pth if pth.is_absolute() else config.work_dir / pth
185
+ return get_path(pth, base=config.input_dir)
111
186
 
112
187
 
113
188
  def get_output_path(pth):
@@ -129,6 +204,4 @@ def get_output_path(pth):
129
204
  :group: foxes.config
130
205
 
131
206
  """
132
- if not isinstance(pth, Path):
133
- pth = Path(pth)
134
- return pth if pth.is_absolute() else config.out_dir / pth
207
+ return get_path(pth, base=config.output_dir)
foxes/constants.py CHANGED
@@ -189,7 +189,18 @@ WORK_DIR = "work_dir"
189
189
  :group: foxes.constants
190
190
  """
191
191
 
192
- OUT_DIR = "out_dir"
192
+ INPUT_DIR = "in_dir"
193
+ """Identifier for the input base directory
194
+ :group: foxes.constants
195
+ """
196
+
197
+ OUTPUT_DIR = "out_dir"
193
198
  """Identifier for the default output directory
194
199
  :group: foxes.constants
195
200
  """
201
+
202
+
203
+ NC_ENGINE = "nc_engine"
204
+ """Identifier for the NetCDF engine
205
+ :group: foxes.constants
206
+ """
foxes/core/algorithm.py CHANGED
@@ -74,14 +74,19 @@ class Algorithm(Model):
74
74
  )
75
75
  elif "engine" in engine_pars:
76
76
  engine_pars["engine_type"] = engine_pars.pop("engine")
77
- v = engine_pars.pop("verbosity", verbosity)
78
- try:
79
- e = Engine.new(verbosity=v, **engine_pars)
80
- except TypeError as e:
81
- print(f"\nError while interpreting engine_pars {engine_pars}\n")
82
- raise e
83
- self.print(f"Algorithm '{self.name}': Selecting engine '{e}'")
84
- e.initialize()
77
+
78
+ if "engine_type" in engine_pars:
79
+ try:
80
+ e = Engine.new(verbosity=verbosity, **engine_pars)
81
+ except TypeError as e:
82
+ print(f"\nError while interpreting engine_pars {engine_pars}\n")
83
+ raise e
84
+ self.print(f"Algorithm '{self.name}': Selecting engine '{e}'")
85
+ e.initialize()
86
+ else:
87
+ raise KeyError(
88
+ f"{self.name}: Found unsupported parameters {list(engine_pars.keys())}"
89
+ )
85
90
 
86
91
  @property
87
92
  def farm(self):
foxes/core/engine.py CHANGED
@@ -157,6 +157,36 @@ class Engine(ABC):
157
157
  if self.verbosity >= level:
158
158
  print(*args, **kwargs)
159
159
 
160
+ def map(
161
+ self,
162
+ func,
163
+ inputs,
164
+ *args,
165
+ **kwargs,
166
+ ):
167
+ """
168
+ Runs a function on a list of files
169
+
170
+ Parameters
171
+ ----------
172
+ func: Callable
173
+ Function to be called on each file,
174
+ func(input, *args, **kwargs) -> data
175
+ inputs: array-like
176
+ The input data list
177
+ args: tuple, optional
178
+ Arguments for func
179
+ kwargs: dict, optional
180
+ Keyword arguments for func
181
+
182
+ Returns
183
+ -------
184
+ results: list
185
+ The list of results
186
+
187
+ """
188
+ raise NotImplementedError
189
+
160
190
  @property
161
191
  def loop_dims(self):
162
192
  """
@@ -176,6 +176,39 @@ class FarmController(FarmDataModel):
176
176
 
177
177
  return [m.name for m in tmodels], tmsels
178
178
 
179
+ def find_turbine_types(self, algo):
180
+ """
181
+ Collects the turbine types.
182
+
183
+ Parameters
184
+ ----------
185
+ algo: foxes.core.Algorithm
186
+ The algorithm
187
+
188
+ """
189
+
190
+ # check turbine models, and find turbine types and pre/post-rotor models:
191
+ self.turbine_types = [None for t in algo.farm.turbines]
192
+ for ti, t in enumerate(algo.farm.turbines):
193
+ for mname in t.models:
194
+ if mname in algo.mbook.turbine_types:
195
+ m = algo.mbook.turbine_types[mname]
196
+ if not isinstance(m, TurbineType):
197
+ raise TypeError(
198
+ f"Model {mname} type {type(m).__name__} is not derived from {TurbineType.__name__}"
199
+ )
200
+ if self.turbine_types[ti] is not None:
201
+ raise TypeError(
202
+ f"Two turbine type models found for turbine {ti}: {self.turbine_types[ti].name} and {mname}"
203
+ )
204
+ m.name = mname
205
+ self.turbine_types[ti] = m
206
+
207
+ if self.turbine_types[ti] is None:
208
+ raise ValueError(
209
+ f"Turbine {ti}, {t.name}: Missing a turbine type model among models {t.models}"
210
+ )
211
+
179
212
  def collect_models(self, algo):
180
213
  """
181
214
  Analyze and gather turbine models, based on the
@@ -188,22 +221,18 @@ class FarmController(FarmDataModel):
188
221
 
189
222
  """
190
223
 
224
+ if self.turbine_types is None:
225
+ self.find_turbine_types(algo)
226
+
191
227
  # check turbine models, and find turbine types and pre/post-rotor models:
192
- self.turbine_types = [None for t in algo.farm.turbines]
193
228
  prer_models = [[] for t in algo.farm.turbines]
194
229
  postr_models = [[] for t in algo.farm.turbines]
230
+ ttypes = {m.name: m for m in self.turbine_types}
195
231
  for ti, t in enumerate(algo.farm.turbines):
196
232
  prer = None
197
233
  for mi, mname in enumerate(t.models):
198
- istype = False
199
- if mname in algo.mbook.turbine_types:
200
- m = algo.mbook.turbine_types[mname]
201
- if not isinstance(m, TurbineType):
202
- raise TypeError(
203
- f"Model {mname} type {type(m).__name__} is not derived from {TurbineType.__name__}"
204
- )
205
- models = [m]
206
- istype = True
234
+ if mname in ttypes:
235
+ models = [ttypes[mname]]
207
236
  elif mname in algo.mbook.turbine_models:
208
237
  m = algo.mbook.turbine_models[mname]
209
238
  models = m.models if isinstance(m, FarmDataModelList) else [m]
@@ -212,21 +241,14 @@ class FarmController(FarmDataModel):
212
241
  raise TypeError(
213
242
  f"Model {mname} type {type(mm).__name__} is not derived from {TurbineModel.__name__}"
214
243
  )
244
+ m.name = mname
215
245
  else:
216
246
  raise KeyError(
217
247
  f"Model {mname} not found in model book types or models"
218
248
  )
219
249
 
220
- if istype:
221
- if self.turbine_types[ti] is None:
222
- self.turbine_types[ti] = m
223
- else:
224
- raise ValueError(
225
- f"Turbine {ti}, {t.name}: Multiple turbine types found in self.turbine_models list, {self.turbine_types[ti].name} and {mname}"
226
- )
227
-
250
+ prer = None
228
251
  for m in models:
229
- m.name = mname
230
252
  if prer is None:
231
253
  prer = m.pre_rotor
232
254
  elif not prer and m.pre_rotor:
@@ -238,11 +260,6 @@ class FarmController(FarmDataModel):
238
260
  else:
239
261
  postr_models[ti].append(m)
240
262
 
241
- if self.turbine_types[ti] is None:
242
- raise ValueError(
243
- f"Turbine {ti}, {t.name}: Missing a turbine type model among models {t.models}"
244
- )
245
-
246
263
  # analyze models:
247
264
  mnames_pre, tmsels_pre = self._analyze_models(
248
265
  algo, pre_rotor=True, models=prer_models
foxes/core/states.py CHANGED
@@ -277,7 +277,7 @@ class ExtendedStates(States):
277
277
  The output variable names
278
278
 
279
279
  """
280
- return self.states.output_point_vars(algo)
280
+ return self.pmodels.output_point_vars(algo)
281
281
 
282
282
  def calculate(self, algo, mdata, fdata, tdata):
283
283
  """
foxes/core/wind_farm.py CHANGED
@@ -1,3 +1,8 @@
1
+ import numpy as np
2
+
3
+ from foxes.config import config
4
+
5
+
1
6
  class WindFarm:
2
7
  """
3
8
  The wind farm.
@@ -76,3 +81,107 @@ class WindFarm:
76
81
 
77
82
  """
78
83
  return [t.name for t in self.turbines]
84
+
85
+ @property
86
+ def xy_array(self):
87
+ """
88
+ Returns an array of the wind farm ground points
89
+
90
+ Returns
91
+ -------
92
+ xya: numpy.ndarray
93
+ The turbine ground positions, shape: (n_turbines, 2)
94
+
95
+ """
96
+ return np.array([t.xy for t in self.turbines], dtype=config.dtype_double)
97
+
98
+ def get_xy_bounds(self, extra_space=None, algo=None):
99
+ """
100
+ Returns min max points of the wind farm ground points
101
+
102
+ Parameters
103
+ ----------
104
+ extra_space: float or str, optional
105
+ The extra space, either float in m,
106
+ or str for units of D, e.g. '2.5D'
107
+ algo: foxes.core.Algorithm, optional
108
+ The algorithm
109
+
110
+ Returns
111
+ -------
112
+ x_mima: numpy.ndarray
113
+ The (x_min, x_max) point
114
+ y_mima: numpy.ndarray
115
+ The (y_min, y_max) point
116
+
117
+ """
118
+ if self.boundary is not None:
119
+ xy = None
120
+ p_min, p_max = self.boundary.p_min(), self.boundary.p_max()
121
+ else:
122
+ xy = self.xy_array
123
+ p_min, p_max = np.min(xy, axis=0), np.max(xy, axis=0)
124
+
125
+ if extra_space is not None:
126
+ if isinstance(extra_space, str):
127
+ assert (
128
+ algo is not None
129
+ ), f"WindFarm: require algo argument for extra_space '{extra_space}'"
130
+ assert (
131
+ len(extra_space) > 1 and extra_space[-1] == "D"
132
+ ), f"Expecting float or str like '2.5D', got extra_space = '{extra_space}'"
133
+ extra_space = float(extra_space[:-1])
134
+ rds = self.get_rotor_diameters(algo)
135
+ if xy is None:
136
+ extra_space *= np.max(rds)
137
+ else:
138
+ p_min = np.min(xy - extra_space * rds[:, None], axis=0)
139
+ p_max = np.max(xy + extra_space * rds[:, None], axis=0)
140
+ return p_min, p_max
141
+
142
+ p_min -= extra_space
143
+ p_max += extra_space
144
+
145
+ return p_min, p_max
146
+
147
+ def get_rotor_diameters(self, algo):
148
+ """
149
+ Gets the rotor diameters
150
+
151
+ Parameters
152
+ ----------
153
+ algo: foxes.core.Algorithm
154
+ The algorithm
155
+
156
+ Returns
157
+ -------
158
+ rds: numpy.ndarray
159
+ The rotor diameters, shape: (n_turbienes,)
160
+
161
+ """
162
+ rds = [
163
+ t.D if t.D is not None else algo.farm_controller.turbine_types[i].D
164
+ for i, t in enumerate(self.turbines)
165
+ ]
166
+ return np.array(rds, dtype=config.dtype_double)
167
+
168
+ def get_hub_heights(self, algo):
169
+ """
170
+ Gets the hub heights
171
+
172
+ Parameters
173
+ ----------
174
+ algo: foxes.core.Algorithm
175
+ The algorithm
176
+
177
+ Returns
178
+ -------
179
+ hhs: numpy.ndarray
180
+ The hub heights, shape: (n_turbienes,)
181
+
182
+ """
183
+ hhs = [
184
+ t.H if t.H is not None else algo.farm_controller.turbine_types[i].H
185
+ for i, t in enumerate(self.turbines)
186
+ ]
187
+ return np.array(hhs, dtype=config.dtype_double)
foxes/engines/dask.py CHANGED
@@ -22,9 +22,11 @@ def load_dask():
22
22
  """On-demand loading of the dask package"""
23
23
  global dask, ProgressBar, delayed
24
24
  if dask is None:
25
- dask = import_module("dask", hint="pip install dask")
25
+ dask = import_module("dask")
26
26
  ProgressBar = import_module(
27
- "dask.diagnostics", hint="pip install dask"
27
+ "dask.diagnostics",
28
+ pip_hint="pip install dask",
29
+ conda_hint="conda install dask -c conda-forge",
28
30
  ).ProgressBar
29
31
  delayed = dask.delayed
30
32
 
@@ -33,7 +35,7 @@ def load_distributed():
33
35
  """On-demand loading of the distributed package"""
34
36
  global distributed
35
37
  if distributed is None:
36
- distributed = import_module("distributed", hint="pip install distributed")
38
+ distributed = import_module("distributed")
37
39
 
38
40
 
39
41
  class DaskBaseEngine(Engine):
@@ -98,6 +100,49 @@ class DaskBaseEngine(Engine):
98
100
  dask.config.set(**self.dask_config)
99
101
  super().initialize()
100
102
 
103
+ def map(
104
+ self,
105
+ func,
106
+ inputs,
107
+ *args,
108
+ **kwargs,
109
+ ):
110
+ """
111
+ Runs a function on a list of files
112
+
113
+ Parameters
114
+ ----------
115
+ func: Callable
116
+ Function to be called on each file,
117
+ func(input, *args, **kwargs) -> data
118
+ inputs: array-like
119
+ The input data list
120
+ args: tuple, optional
121
+ Arguments for func
122
+ kwargs: dict, optional
123
+ Keyword arguments for func
124
+
125
+ Returns
126
+ -------
127
+ results: list
128
+ The list of results
129
+
130
+ """
131
+ if len(inputs) == 0:
132
+ return []
133
+ elif len(inputs) == 1:
134
+ return [func(inputs[0], *args, **kwargs)]
135
+ else:
136
+ inptl = np.array_split(inputs, min(self.n_procs, len(inputs)))
137
+ jobs = []
138
+ for subi in inptl:
139
+ jobs.append(_run_map(func, subi, *args, **kwargs))
140
+ results = dask.compute(jobs)[0]
141
+ out = []
142
+ for r in results:
143
+ out += r
144
+ return out
145
+
101
146
  def chunk_data(self, data):
102
147
  """
103
148
  Applies the selected chunking
@@ -280,6 +325,36 @@ class XArrayEngine(DaskBaseEngine):
280
325
 
281
326
  """
282
327
 
328
+ def map(
329
+ self,
330
+ func,
331
+ inputs,
332
+ *args,
333
+ **kwargs,
334
+ ):
335
+ """
336
+ Runs a function on a list of files
337
+
338
+ Parameters
339
+ ----------
340
+ func: Callable
341
+ Function to be called on each file,
342
+ func(input, *args, **kwargs) -> data
343
+ inputs: array-like
344
+ The input data list
345
+ args: tuple, optional
346
+ Arguments for func
347
+ kwargs: dict, optional
348
+ Keyword arguments for func
349
+
350
+ Returns
351
+ -------
352
+ results: list
353
+ The list of results
354
+
355
+ """
356
+ return [func(input, *args, **kwargs) for input in inputs]
357
+
283
358
  def run_calculation(
284
359
  self,
285
360
  algo,
@@ -476,6 +551,12 @@ def _run_lazy(algo, model, iterative, chunk_store, i0_t0, *data, **cpars):
476
551
  return results, cstore
477
552
 
478
553
 
554
+ @delayed
555
+ def _run_map(func, inputs, *args, **kwargs):
556
+ """Helper function for running map func on proc"""
557
+ return [func(x, *args, **kwargs) for x in inputs]
558
+
559
+
479
560
  class DaskEngine(DaskBaseEngine):
480
561
  """
481
562
  The dask engine for delayed foxes calculations.
@@ -969,8 +1050,11 @@ class SlurmClusterEngine(LocalClusterEngine):
969
1050
  nodes = cargs.pop("nodes", 1)
970
1051
 
971
1052
  dask_jobqueue = import_module(
972
- "dask_jobqueue", hint="pip install setuptools dask-jobqueue"
1053
+ "dask_jobqueue",
1054
+ pip_hint="pip install setuptools dask-jobqueue",
1055
+ conda_hint="conda install setuptools dask-jobqueue -c conda-forge",
973
1056
  )
1057
+
974
1058
  self._cluster = dask_jobqueue.SLURMCluster(**cargs)
975
1059
  self._cluster.scale(jobs=nodes)
976
1060
  self._cluster = self._cluster.__enter__()