From d021afd186b7efa2ff48a10e63ea3c5a107a48d5 Mon Sep 17 00:00:00 2001 From: Pranjal Dhole Date: Wed, 13 Nov 2024 20:33:13 +0100 Subject: [PATCH] harmonization of simulation solver method imports --- causal_inference/base/lv_simulator.py | 45 +++++---------------- causal_inference/base/lv_system.py | 33 ++++++++++----- causal_inference/base/ode_solver.py | 31 ++++++++------ causal_inference/base/runge_kutta_solver.py | 16 ++++---- causal_inference/config.py | 2 +- causal_inference/utils/visualisations.py | 21 ++++++++++ causal_inference/utils/writer.py | 12 ++++++ 7 files changed, 94 insertions(+), 66 deletions(-) create mode 100644 causal_inference/utils/visualisations.py create mode 100644 causal_inference/utils/writer.py diff --git a/causal_inference/base/lv_simulator.py b/causal_inference/base/lv_simulator.py index d9bce29..d333f02 100644 --- a/causal_inference/base/lv_simulator.py +++ b/causal_inference/base/lv_simulator.py @@ -6,37 +6,12 @@ import logging import datetime -import h5py -import matplotlib.pyplot as plt - from causal_inference.config import RESULTS_DIR from causal_inference.utils.log_config import log_LV_params from causal_inference.base.ode_solver import ODE_solver from causal_inference.base.runge_kutta_solver import RungeKuttaSolver - -def _save_population(prey_list, predator_list): - filename = os.path.join(RESULTS_DIR, 'populations.h5') - hf = h5py.File(filename, 'w') - hf.create_dataset('prey_pop', data=prey_list) - hf.create_dataset('pred_pop', data=predator_list) - hf.close() - -def plot_population_over_time(prey_list, predator_list, save=True, filename='predator_prey'): - fig = plt.figure(figsize=(15, 5)) - ax = fig.add_subplot(2, 1, 1) - PreyLine, = plt.plot(prey_list , color='g') - PredatorsLine, = plt.plot(predator_list, color='r') - ax.set_xscale('log') - - plt.legend([PreyLine, PredatorsLine], ['Prey', 'Predators']) - plt.ylabel('Population') - plt.xlabel('Time') - if save: - plt.savefig(os.path.join(RESULTS_DIR, f"{filename}.svg"), - format='svg', transparent=False, bbox_inches='tight') - else: - plt.show() - plt.close() +from causal_inference.utils.writer import _save_population +from causal_inference.utils.visualisations import plot_population_over_time def get_solver(method): ''' @@ -50,15 +25,15 @@ def get_solver(method): raise AssertionError(f'{method} is not implemented!') return solver -def main(method): +def main(method, results_dir): ''' Main function that solves LV system. ''' log_LV_params() solver = get_solver(method) prey_list, predator_list = solver._solve() - _save_population(prey_list, predator_list) - plot_population_over_time(prey_list, predator_list) + _save_population(prey_list, predator_list, solver.time_stamps, results_dir) + plot_population_over_time(prey_list, predator_list, solver.time_stamps, results_dir) if __name__ == '__main__': PARSER = argparse.ArgumentParser() @@ -68,12 +43,12 @@ def main(method): choices=['RK4', 'ODE'], default='RK4') ARGS = PARSER.parse_args() - RESULTS_DIR = os.path.join(RESULTS_DIR, '{}_{}'.format(datetime.datetime.now().strftime("%Y%h%d_%H_%M_%S"), str(ARGS.outdir))) + results_dir = os.path.join(RESULTS_DIR, '{}_{}'.format(datetime.datetime.now().strftime("%Y%h%d_%H_%M_%S"), str(ARGS.outdir))) - if not os.path.exists(RESULTS_DIR): - os.makedirs(RESULTS_DIR) + if not os.path.exists(results_dir): + os.makedirs(results_dir) - LOG_FILE = os.path.join(RESULTS_DIR, f"{ARGS.logfile}.txt") # write logg to this file + LOG_FILE = os.path.join(results_dir, f"{ARGS.logfile}.txt") # write logg to this file logging.basicConfig( level=logging.INFO, handlers=[ @@ -82,4 +57,4 @@ def main(method): ] ) solver = ARGS.solver - main(solver) + main(solver, results_dir) diff --git a/causal_inference/base/lv_system.py b/causal_inference/base/lv_system.py index b580e78..af3291c 100644 --- a/causal_inference/base/lv_system.py +++ b/causal_inference/base/lv_system.py @@ -1,20 +1,31 @@ #!/usr/bin/python # -*- coding: utf-8 -*- +import numpy as np from causal_inference.config import LV_PARAMS class LotkaVolterra(): ''' - Class simulates predator-prey dynamics and solves it with 4th order Runge-Kutta method. + Base Lotka-Volterra Class that defines a predator-prey system. ''' - def __init__(self): - self.A = LV_PARAMS['A'] - self.B = LV_PARAMS['B'] - self.C = LV_PARAMS['C'] - self.D = LV_PARAMS['D'] - self.time = LV_PARAMS['INITIAL_TIME'] - self.step_size = LV_PARAMS['STEP_SIZE'] - self.max_iterations = LV_PARAMS['MAX_ITERATIONS'] + def __init__(self, + A=LV_PARAMS['A'], B=LV_PARAMS['B'], C=LV_PARAMS['C'], D=LV_PARAMS['D'], + prey_population=LV_PARAMS['INITIAL_PREY_POPULATION'], + pred_population=LV_PARAMS['INITIAL_PREDATOR_POPULATION'], + total_time=LV_PARAMS['TOTAL_TIME'], step_size=LV_PARAMS['STEP_SIZE'], + max_iter=LV_PARAMS['MAX_ITERATIONS']): + # Lotka-Volterra parameters + self.A = A + self.B = B + self.C = C + self.D = D - self.prey_population = LV_PARAMS['INITIAL_PREY_POPULATION'] - self.predator_population = LV_PARAMS['INITIAL_PREDATOR_POPULATION'] \ No newline at end of file + self.prey_population = prey_population # Initial prey population + self.predator_population = pred_population # Initial predator population + + self.init_time = 0 # initial time + self.total_time = total_time # total time in units + self.step_size = step_size # increment for each time step + self.max_iterations = max_iter # tolerance parameter + + self.time_stamps = np.arange(self.init_time, self.total_time, self.step_size) diff --git a/causal_inference/base/ode_solver.py b/causal_inference/base/ode_solver.py index 7ddb0e1..7f47151 100644 --- a/causal_inference/base/ode_solver.py +++ b/causal_inference/base/ode_solver.py @@ -5,7 +5,6 @@ import numpy as np from scipy import integrate -from causal_inference.config import LV_PARAMS from causal_inference.base.lv_system import LotkaVolterra class ODE_solver(LotkaVolterra): @@ -14,20 +13,28 @@ class ODE_solver(LotkaVolterra): ''' def __init__(self): super().__init__() - logging.info('Solving Lotka-Volterra predator-prey dynamics odeint solver') + logging.info('Simulating Lotka-Volterra predator-prey dynamics with odeint solver') @staticmethod - def LV_derivative(X, t, alpha, beta, delta, gamma): - x, y = X - dotx = x * (alpha - beta * y) - doty = y * (-delta + gamma * x) + def LV_derivative(t, Z, A, B, C, D): + ''' + Returns the rate of change of predator and prey population + ''' + x, y = Z + dotx = x * (A - B * y) + doty = y * (-C + D * x) return np.array([dotx, doty]) def _solve(self): - logging.info('Computing population over time...') - t = np.arange(0.,self.max_iterations, self.step_size) - X0 = [self.prey_population, self.predator_population] - res = integrate.odeint(self.LV_derivative, X0, t, args=(self.A, self.B, self.C, self.D)) - prey_list, predator_list = res.T + ''' + ODE solver that returns the predator and prey populations at each time step in time series. + ''' + logging.info(f'Computing population over {self.total_time} generation with step size of {self.step_size}...') + + INIT_POP = [self.prey_population, self.predator_population] + sol = integrate.solve_ivp(self.LV_derivative, [self.init_time, self.total_time], INIT_POP, args=(self.A, self.B, self.C, self.D), dense_output=True) + prey_list, predator_list = sol.sol(self.time_stamps) + logging.info('done!') - return prey_list, predator_list \ No newline at end of file + + return prey_list, predator_list diff --git a/causal_inference/base/runge_kutta_solver.py b/causal_inference/base/runge_kutta_solver.py index 8a54c89..440d733 100644 --- a/causal_inference/base/runge_kutta_solver.py +++ b/causal_inference/base/runge_kutta_solver.py @@ -1,7 +1,6 @@ #!/usr/bin/python # -*- coding: utf-8 -*- import logging -from math import ceil from causal_inference.base.lv_system import LotkaVolterra class RungeKuttaSolver(LotkaVolterra): @@ -11,7 +10,6 @@ class RungeKuttaSolver(LotkaVolterra): def __init__(self): super().__init__() logging.info('Solving Lotka-Volterra predator-prey dynamics with 4th order Runge-Kutta method') - self.time_stamp = [self.time] self.prey_list = [self.prey_population] self.predator_list = [self.predator_population] @@ -22,8 +20,6 @@ def compute_predator_rate(self, current_prey, current_predators): return - self.C * current_predators + self.D * current_prey * current_predators def runge_kutta_update(self, current_prey, current_predators): - self.time = self.time + self.step_size - self.time_stamp.append(self.time) k1_prey = self.step_size * self.compute_prey_rate(current_prey, current_predators) k1_pred = self.step_size * self.compute_predator_rate(current_prey, current_predators) @@ -46,11 +42,17 @@ def runge_kutta_update(self, current_prey, current_predators): return new_prey_population, new_predator_population def _solve(self): + ''' + Runge-Kutta solver that returns the predator and prey populations at each time step in time series. + ''' + #initial population current_prey, current_predators = self.prey_population, self.predator_population - logging.info('Computing population over time...') - for gen_idx in range(ceil(self.max_iterations/self.step_size)): + + logging.info(f'Computing population over {self.total_time} generation with step size of {self.step_size}...') + + for step_idx in self.time_stamps[1:]: current_prey, current_predators = self.runge_kutta_update(current_prey, current_predators) - msg= f'Gen: {gen_idx} | Prey population: {current_prey} | Predator population: {current_predators}' + msg= f'Step: {step_idx} | Prey population: {current_prey} | Predator population: {current_predators}' logging.info(msg) print('Done!') return self.prey_list, self.predator_list \ No newline at end of file diff --git a/causal_inference/config.py b/causal_inference/config.py index 4acd3de..bf940b2 100644 --- a/causal_inference/config.py +++ b/causal_inference/config.py @@ -12,7 +12,7 @@ 'C' : 3.0, 'D' : 5.0, 'STEP_SIZE' : 0.01, -'INITIAL_TIME' : 0, +'TOTAL_TIME' : 20, 'INITIAL_PREY_POPULATION' : 60, 'INITIAL_PREDATOR_POPULATION' : 25, 'MAX_ITERATIONS' : 200 diff --git a/causal_inference/utils/visualisations.py b/causal_inference/utils/visualisations.py new file mode 100644 index 0000000..220ba98 --- /dev/null +++ b/causal_inference/utils/visualisations.py @@ -0,0 +1,21 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +from os.path import join +import matplotlib.pyplot as plt + +def plot_population_over_time(prey_list, predator_list, time_stamps, results_dir, save=True, filename='predator_prey'): + fig = plt.figure(figsize=(15, 5)) + ax = fig.add_subplot(2, 1, 1) + PreyLine, = plt.plot(time_stamps, prey_list, color='g') + PredatorsLine, = plt.plot(time_stamps, predator_list, color='r') + ax.set_xscale('log') + + plt.legend([PreyLine, PredatorsLine], ['Prey', 'Predators']) + plt.ylabel('Population') + plt.xlabel('Time') + if save: + plt.savefig(join(results_dir, f"{filename}.svg"), + format='svg', transparent=False, bbox_inches='tight') + else: + plt.show() + plt.close() \ No newline at end of file diff --git a/causal_inference/utils/writer.py b/causal_inference/utils/writer.py new file mode 100644 index 0000000..2d5da4e --- /dev/null +++ b/causal_inference/utils/writer.py @@ -0,0 +1,12 @@ +#!/usr/bin/python +# -*- coding: utf-8 -*- +from os.path import join +import h5py + +def _save_population(prey_list, predator_list, time_stamps, results_dir): + filename = join(results_dir, 'populations.h5') + hf = h5py.File(filename, 'w') + hf.create_dataset('time_stamp', data=time_stamps) + hf.create_dataset('prey_pop', data=prey_list) + hf.create_dataset('pred_pop', data=predator_list) + hf.close()