Merge pull request #77 from EcohydrologyTeam/performance-pre-init

Pre-Initialize model xarray with timesteps

.gitignore

@@ -84,6 +84,9 @@ target/
 .ipynb_checkpoints
 .virtual_documents
 
+# profiling
+*.prof
+
 # IPython
 profile_default/
 ipython_config.py

examples/dev_sandbox/performance_profiling_tsm.py

@@ -0,0 +1,106 @@
+import clearwater_modules
+import time
+import sys
+import logging
+import numpy as np
+import xarray as xr
+import cProfile
+
+from typing import (
+    List
+)
+
+def create_data_array(
+    gridsize: int,
+    value: int | float  
+):
+    data = np.full(gridsize, value)
+    data_array = xr.DataArray(data)
+    return data_array
+
+def run_performance_test(
+        iterations_list: List[int],
+        gridsize_list: List[int],
+        log_file: str,
+        detailed_profile: bool
+):
+    """Log the performance for different numbers of iterations."""
+    # set up logger
+    log_format = '%(asctime)s,%(levelname)s,%(message)s'
+    logging.basicConfig(
+        filename=log_file,
+        level=logging.INFO,
+        format=log_format,
+    )
+    if detailed_profile:
+        logging.info(f"timestep,MB,time_per_iter")
+    else:
+        logging.info(f"iters,gridsize,avg_increment_time,total_increment_time")
+
+    for iters in iterations_list:
+        for gridsize in gridsize_list:
+            # define starting state values
+            if gridsize == 1:
+                state_i = {
+                    'water_temp_c': 40.0,
+                    'surface_area': 1.0,
+                    'volume': 1.0,
+                }
+                meteo_parameters = {'wind_c': 1.0}
+
+            else:
+                state_i = {
+                    'water_temp_c': create_data_array(gridsize, 40.0),
+                    'surface_area': create_data_array(gridsize, 1.0),
+                    'volume': create_data_array(gridsize, 1.0),
+                }
+                meteo_parameters = {
+                    'wind_c': 1.0
+                }
+
+            ti = time.time()
+
+            # instantiate the TSM module
+            tsm = clearwater_modules.tsm.EnergyBudget(
+                time_steps=iters,
+                initial_state_values=state_i,
+                meteo_parameters=meteo_parameters,
+                track_dynamic_variables=False,
+            )
+
+            t2 = time.time()
+            curr_time = t2
+            for _ in range(iters):
+                tsm.increment_timestep()
+
+                # detailed profiling
+                if detailed_profile:
+                    if (_ % 100 == 0) and (_ != 0):
+                        current_time = (time.time() - curr_time) / 100
+                        logging.info(f"{_},{tsm.dataset.nbytes * 0.000001}, {current_time}")
+                        curr_time = time.time()
+
+            # wrap up
+            avg_increment_time = (time.time() - t2) / iters
+            total_run_time = time.time() - ti
+            if not detailed_profile:
+                logging.info(f"{iters},{gridsize},{avg_increment_time},{total_run_time}")
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        print("Usage: python performance_profiling_tsm.py <log_file>")
+        sys.exit(1)
+
+    log_file = f"{sys.argv[1]}.log"
+    cprofile_file = f"{sys.argv[1]}.prof"
+    iterations_list = [1, 10, 100, 1000, 10000, 100000]
+    gridsize_list = [1, 1000, 10000]
+    detailed_profile = False
+    # iterations_list = [1000]
+    # gridsize_list = [1]
+    # detailed_profile = True
+    # profiler = cProfile.Profile()
+    # profiler.enable()
+    run_performance_test(iterations_list, gridsize_list, log_file, detailed_profile)
+    # profiler.disable()
+    # profiler.dump_stats(cprofile_file)

examples/dev_sandbox/prof.py

@@ -2,23 +2,53 @@
 import clearwater_modules
 import time 
 import sys
+import xarray as xr
+import numpy as np
 
-def main(iters: int):
+
+def main(iters: int, type: str):
     ti = time.time()
     # define starting state values
-    state_i = {
-        'water_temp_c': 40.0,
-        'surface_area': 1.0,
-        'volume': 1.0,
-    }
+    if type == 'baseline':
+        state_i = {
+            'water_temp_c': 40.0,
+            'surface_area': 1.0,
+            'volume': 1.0,
+        }
+    elif type in ['arrays', 'hotstart']:
+        state_i = {
+            'water_temp_c': xr.DataArray(
+                np.full(10, 40),
+                dims='cell',
+                coords={'cell': np.arange(10)}),
+            'surface_area': xr.DataArray(
+                np.full(10, 1.0),
+                dims='cell',
+                coords={'cell': np.arange(10)}),
+            'volume': xr.DataArray(
+                np.full(10, 1.0),
+                dims='cell',
+                coords={'cell': np.arange(10)}),
+        }
+
 
     # instantiate the TSM module
     tsm = clearwater_modules.tsm.EnergyBudget(
+        time_steps=iters,
         initial_state_values=state_i,
         meteo_parameters={'wind_c': 1.0},
+        updateable_static_variables=['wind_c']
     )
-    print(tsm.static_variable_values)
+
     t2 = time.time()
+
+    if type == 'hotstart':
+        tsm = clearwater_modules.tsm.EnergyBudget(
+            time_steps=iters,
+            hotstart_dataset=tsm.dataset,
+        )
+        t2 = time.time()
+
     for _ in range(iters):
         tsm.increment_timestep()
     print(f'Increment timestep speed (average of {iters}): {(time.time() - t2) / 100}')
@@ -35,4 +65,4 @@ def main(iters: int):
         print('No argument given, defaulting to 100 iteration.')
         iters = 100
 
-    main(iters=iters)
+    main(iters=iters, type='baseline')