diff --git a/src/patatune/mopso/mopso.py b/src/patatune/mopso/mopso.py index d4c12f5..4db2ffd 100644 --- a/src/patatune/mopso/mopso.py +++ b/src/patatune/mopso/mopso.py @@ -263,19 +263,41 @@ def step(self, max_iterations_without_improvement=None): [particle.position for particle in self.particles]) [particle.set_fitness(optimization_output[p_id]) for p_id, particle in enumerate(self.particles)] - FileManager.save_csv([np.concatenate([particle.position, np.ravel( - particle.fitness * self.objective.directions)]) for particle in self.particles], - 'history/iteration' + str(self.iteration) + '.csv', - headers=self.param_names + self.objective.objective_names) - self.history[self.iteration] = np.array( - [(particle.id, particle.position, particle.fitness * self.objective.directions) for particle in self.particles], - dtype=np.dtype([('id', int), ('position', float, (self.num_params,)), ('fitness', float, (self.objective.num_objectives,))]) - ) + + if FileManager.saving_history_enabled: + FileManager.save_csv([np.concatenate([particle.position, np.ravel( + particle.fitness * self.objective.directions)]) for particle in self.particles], + 'history/iteration' + str(self.iteration) + '.csv', + headers=self.param_names + self.objective.objective_names) + + # Store particle history as list of dictionaries for flexibility + particle_data = [] + for particle in self.particles: + particle_data.append({ + 'id': particle.id, + 'position': np.array(particle.position), + 'fitness': np.array(particle.fitness * self.objective.directions, dtype=float) + }) + self.history[self.iteration] = particle_data + crowding_distances = self.update_pareto_front() - self.history['pareto_front'] = np.array( - [(particle.position, particle.fitness * self.objective.directions) for particle in self.pareto_front], - dtype=[('position', float, (self.num_params,)), ('fitness', float, (self.objective.num_objectives,))] - ) + + if FileManager.saving_history_enabled: + # Save pareto front history per iteration (not overwriting) + pareto_key = f'pareto_front_{self.iteration}' + pareto_data = [] + for particle in self.pareto_front: + pareto_data.append({ + 'position': np.array(particle.position), + 'fitness': np.array(particle.fitness * self.objective.directions, dtype=float) + }) + self.history[pareto_key] = pareto_data + + # Export pareto front to CSV per iteration + FileManager.save_csv([np.concatenate([particle.position, np.ravel( + particle.fitness * self.objective.directions)]) for particle in self.pareto_front], + 'history/pareto_iteration' + str(self.iteration) + '.csv', + headers=self.param_names + self.objective.objective_names) for particle in self.particles: particle.update_velocity(self.pareto_front, crowding_distances, diff --git a/src/patatune/util.py b/src/patatune/util.py index a7c2201..b913c2f 100644 --- a/src/patatune/util.py +++ b/src/patatune/util.py @@ -137,6 +137,7 @@ class FileManager: saving_json_enabled (bool): Flag to enable/disable JSON saving. saving_zarr_enabled (bool): Flag to enable/disable Zarr saving. saving_pickle_enabled (bool): Flag to enable/disable Pickle saving. + saving_history_enabled (bool): Flag to enable/disable history saving (per-iteration particle and Pareto front data). loading_enabled (bool): Global flag to enable/disable loading. headers_enabled (bool): Flag to enable/disable headers when saving/loading CSV files. working_dir (str): Directory where files will be saved/loaded from. @@ -147,6 +148,7 @@ class FileManager: saving_json_enabled = True saving_zarr_enabled = False saving_pickle_enabled = True + saving_history_enabled = True loading_enabled = False headers_enabled = False working_dir = "tmp" @@ -328,7 +330,29 @@ def save_zarr(cls, obj, filename, **kwargs): group_name = key group = root_group.create_group(group_name) - group.create_dataset("data", data=value, overwrite=True) + + # Handle list of dictionaries (new format for particle/pareto history) + if isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict): + # Convert list of dicts to separate arrays for each field + for field_name in value[0].keys(): + field_data = [item[field_name] for item in value] + try: + # Try to stack as numpy array + field_array = np.array(field_data) + # Check if the array has object dtype (which happens with mixed types) + if field_array.dtype == object: + # Use Pickle codec for object arrays + from numcodecs import Pickle + group.create_dataset(field_name, data=field_array, overwrite=True, object_codec=Pickle()) + else: + group.create_dataset(field_name, data=field_array, overwrite=True) + except (ValueError, TypeError): + # If stacking fails, save as object array with Pickle codec + from numcodecs import Pickle + group.create_dataset(field_name, data=np.array(field_data, dtype=object), overwrite=True, object_codec=Pickle()) + else: + # Handle numpy arrays (backward compatibility) + group.create_dataset("data", data=value, overwrite=True) root_group.attrs.update(kwargs) store.close() diff --git a/tests/test_pareto_history.py b/tests/test_pareto_history.py new file mode 100644 index 0000000..0d053a4 --- /dev/null +++ b/tests/test_pareto_history.py @@ -0,0 +1,154 @@ +"""Test pareto front history saving functionality.""" + +import patatune +import numpy as np +import os +import shutil + + +def test_pareto_history_enabled(): + """Test that pareto front history is saved when enabled.""" + # Clean up test directory + test_dir = "tmp/test_pareto_history_enabled" + if os.path.exists(test_dir): + shutil.rmtree(test_dir) + + # Setup + lb = [0.0, 0.0] + ub = [5.0, 3.0] + num_agents = 10 + num_iterations = 3 + + def f(params): + return [[4 * p[0]**2 + 4 * p[1]**2, (p[0] - 5)**2 + (p[1] - 5)**2] for p in params] + + patatune.FileManager.working_dir = test_dir + patatune.FileManager.saving_history_enabled = True + + objective = patatune.Objective([f]) + pso = patatune.MOPSO(objective=objective, lower_bounds=lb, upper_bounds=ub, + num_particles=num_agents, + inertia_weight=0.5, cognitive_coefficient=1, social_coefficient=1) + + # Run optimization + pso.optimize(num_iterations) + + # Verify history keys + assert len(pso.history) == num_iterations * 2 # particles + pareto for each iteration + + # Verify CSV files exist for each iteration + for i in range(num_iterations): + particle_csv = os.path.join(test_dir, f'history/iteration{i}.csv') + pareto_csv = os.path.join(test_dir, f'history/pareto_iteration{i}.csv') + + assert os.path.exists(particle_csv), f"Particle CSV for iteration {i} not found" + assert os.path.exists(pareto_csv), f"Pareto CSV for iteration {i} not found" + + # Verify history dictionary contains data + assert i in pso.history, f"Iteration {i} not in history" + assert f'pareto_front_{i}' in pso.history, f"Pareto front {i} not in history" + + # Verify data structure + assert isinstance(pso.history[i], list), f"Iteration {i} history is not a list" + assert isinstance(pso.history[f'pareto_front_{i}'], list), f"Pareto front {i} history is not a list" + + if len(pso.history[i]) > 0: + assert 'id' in pso.history[i][0], f"Particle data missing 'id' field" + assert 'position' in pso.history[i][0], f"Particle data missing 'position' field" + assert 'fitness' in pso.history[i][0], f"Particle data missing 'fitness' field" + + if len(pso.history[f'pareto_front_{i}']) > 0: + assert 'position' in pso.history[f'pareto_front_{i}'][0], f"Pareto data missing 'position' field" + assert 'fitness' in pso.history[f'pareto_front_{i}'][0], f"Pareto data missing 'fitness' field" + + print("✓ Test passed: Pareto history enabled") + + +def test_pareto_history_disabled(): + """Test that pareto front history is NOT saved when disabled.""" + # Clean up test directory + test_dir = "tmp/test_pareto_history_disabled" + if os.path.exists(test_dir): + shutil.rmtree(test_dir) + + # Setup + lb = [0.0, 0.0] + ub = [5.0, 3.0] + num_agents = 10 + num_iterations = 3 + + def f(params): + return [[4 * p[0]**2 + 4 * p[1]**2, (p[0] - 5)**2 + (p[1] - 5)**2] for p in params] + + patatune.FileManager.working_dir = test_dir + patatune.FileManager.saving_history_enabled = False + + objective = patatune.Objective([f]) + pso = patatune.MOPSO(objective=objective, lower_bounds=lb, upper_bounds=ub, + num_particles=num_agents, + inertia_weight=0.5, cognitive_coefficient=1, social_coefficient=1) + + # Run optimization + pso.optimize(num_iterations) + + # Verify history is empty + assert len(pso.history) == 0, "History should be empty when saving is disabled" + + # Verify CSV files do NOT exist + history_dir = os.path.join(test_dir, 'history') + assert not os.path.exists(history_dir), "History directory should not exist when saving is disabled" + + print("✓ Test passed: Pareto history disabled") + + +def test_pareto_history_growth(): + """Test that pareto front grows/changes over iterations.""" + # Clean up test directory + test_dir = "tmp/test_pareto_history_growth" + if os.path.exists(test_dir): + shutil.rmtree(test_dir) + + # Setup + lb = [0.0, 0.0] + ub = [5.0, 3.0] + num_agents = 20 + num_iterations = 5 + + def f(params): + return [[4 * p[0]**2 + 4 * p[1]**2, (p[0] - 5)**2 + (p[1] - 5)**2] for p in params] + + patatune.FileManager.working_dir = test_dir + patatune.FileManager.saving_history_enabled = True + + objective = patatune.Objective([f]) + pso = patatune.MOPSO(objective=objective, lower_bounds=lb, upper_bounds=ub, + num_particles=num_agents, + inertia_weight=0.5, cognitive_coefficient=1, social_coefficient=1) + + # Run optimization + pso.optimize(num_iterations) + + # Check that pareto front sizes change (typically grow early on) + pareto_sizes = [] + for i in range(num_iterations): + pareto_key = f'pareto_front_{i}' + size = len(pso.history[pareto_key]) + pareto_sizes.append(size) + print(f" Iteration {i}: Pareto front size = {size}") + + # At least some iterations should have non-empty pareto fronts + assert any(size > 0 for size in pareto_sizes), "Pareto front should have solutions" + + print("✓ Test passed: Pareto history growth tracked") + + +if __name__ == "__main__": + print("Testing pareto front history functionality...") + print() + + test_pareto_history_enabled() + test_pareto_history_disabled() + test_pareto_history_growth() + + print() + print("All tests passed! ✓")