Create cprofile_tsm.py

Author: sjordan29

examples/dev_sandbox/cprofile_tsm.py

@@ -0,0 +1,68 @@
+import clearwater_modules
+import numpy as np
+import xarray as xr
+import cProfile
+from pathlib import Path
+
+
+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(
+        iters: int,
+        gridsize: int
+):
+    """Log the performance for different numbers of iterations."""
+
+    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
+        }
+    
+        # 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,
+        )
+
+        for _ in range(iters):
+            tsm.increment_timestep()
+
+
+if __name__ == '__main__':
+    iterations_list = [300000, 400000, 500000] # [1000, 10000, 100000, 150000, 200000]
+    gridsize_list = [10000, 100000]
+    cwd = Path.cwd()
+    print(cwd)
+    root = Path('./examples/dev_sandbox/profiling/')
+    # iterations_list = [1000]
+    # gridsize_list = [1]
+    # detailed_profile = True
+    for iteration in iterations_list:
+        for gridsize in gridsize_list:
+            cprofile_file = root / f'{iteration}_iters_{gridsize}_gridsize.prof'
+            profiler = cProfile.Profile()
+            profiler.enable()
+            run_performance_test(iteration, gridsize)
+            profiler.disable()
+            profiler.dump_stats(cprofile_file)