diff --git a/setup.cfg b/setup.cfg index 7f26022..d6485ae 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = squish -version = 0.1.1 +version = 0.1.2 author = Kenneth Jao author_email = ksjdragon@gmail.com description = squish is Python program which perform simulations for the flow of 'soft' or 'compressible' objects under some energy in a periodic domain. diff --git a/squish/common.py b/squish/common.py index 5aa5ff8..8617966 100644 --- a/squish/common.py +++ b/squish/common.py @@ -15,11 +15,12 @@ STR_TO_ENERGY = { } -def generate_filepath(sim: SimulationMode, fol: Union[str, Path]) -> Path: +def generate_filepath(sim: SimulationMode, fol: Union[str, Path], prec: int = 2) -> Path: energy = sim.energy.title_str width, height = round(sim.domain.w, 2), round(sim.domain.h, 2) - base_path = f"{fol}/{energy}{sim.title_str} - N{sim.domain.n} - {width:.2f}x{height:.2f}" + base_path = f"{fol}/{energy}{sim.title_str} - " + \ + f"N{sim.domain.n} - {width:.{prec}f}x{height:.{prec}f}" i = 1 real_path = Path(base_path) diff --git a/squish/diagram.py b/squish/diagram.py index 71d5d95..827e13f 100644 --- a/squish/diagram.py +++ b/squish/diagram.py @@ -26,17 +26,30 @@ class SimData: def __init__(self, sim: Simulation) -> None: - self.path = sim.path - self.domains = list([DomainParams(s.n, s.w, s.h, s.r) for s in sim]) - self.energies = list([s.energy for s in sim]) - self.voronois = list([s.vor_data for s in sim]) - self.stats = list([s.stats for s in sim]) + if sim is not None: + self.path = sim.path + self.domains = list([DomainParams(s.n, s.w, s.h, s.r) for s in sim]) + self.energies = list([s.energy for s in sim]) + self.voronois = list([s.vor_data for s in sim]) + self.stats = list([s.stats for s in sim]) def __len__(self) -> int: return len(self.domains) + def slice(self, indices: List[int]) -> SimData: + new_data = SimData(None) + new_data.path = self.path + + new_data.domains = list([self.domains[i] for i in indices]) + new_data.energies = list([self.energies[i] for i in indices]) + new_data.voronois = list([self.voronois[i] for i in indices]) + new_data.stats = list([self.stats[i] for i in indices]) + + return new_data + + def hist(self, stat: str, i: int, bins: int = 10, bounds: Optional[Tuple[float, float]] = None, cumul: bool = False, avg: bool = False) -> Tuple[numpy.ndarray, numpy.ndarray]: """Generates a histogram from the selected data. @@ -95,7 +108,7 @@ class Diagram: self.cumulative = cumulative - def generate_frame(self, frame: int, mode: str, fol: str) -> None: + def generate_frame(self, frame: int, mode: str, fol: str, name: str = None) -> None: if mode not in ["save", "open"]: raise ValueError("Not a valid mode for diagrams!") @@ -112,9 +125,11 @@ class Diagram: getattr(self, str(diagram) + '_plot')(frame, axes[it.multi_index]) plt.tight_layout() + if name is None: + name = f"img{frame:05}.png" if mode == "save": - plt.savefig(self.sim.path / fol / f"img{frame:05}.png") + plt.savefig(self.sim.path / fol / name) plt.close(fig) elif mode == "show": plt.show() @@ -290,9 +305,13 @@ class Diagram: (self.sim.path / fol).mkdir(exist_ok=True) combo_list = [] for i in range(cpu_count()): - combo_list.append((self, frames[:int((i+1)*len(frames)/cpu_count())], - fol, len(frames))) + start, end = int(i*len(frames)/cpu_count()), int((i+1)*len(frames)/cpu_count()) + new_dia = Diagram(None, self.diagrams, self.cumulative) + new_dia.sim = self.sim.slice(frames[start:end]) + combo_list.append((new_dia, fol, start, len(frames[start:end]), len(frames))) + # Free up memory, since it's already duplicated to other cores. + self.sim = self.sim.slice([]) with Pool(cpu_count()) as pool: for _ in pool.imap_unordered(render_frame_range, combo_list): pass @@ -300,6 +319,7 @@ class Diagram: print(flush=True) print(f'Wrote to \"{self.sim.path / fol}\".', flush=True) + def render_video(self, time: int, mode: str) -> None: if mode not in ["use_all", "sample"]: raise ValueError("Not a valid mode for videos!") @@ -308,12 +328,11 @@ class Diagram: frames = list(range(len(self.sim))) elif mode == "sample": fps = 30 - if len(self.sim) < fps*time : + if len(self.sim) < fps*time: frames = list(range(len(self.sim))) fps = len(self.sim)/time else: - frames = list(np.round(np.linspace(0, len(self.sim), fps*time)).astype(int)) - print(frames) + frames = list(np.round(np.linspace(0, len(self.sim)-1, fps*time)).astype(int)) self.render_frames(frames, 'temp') path = self.sim.path / 'simulation.mp4' @@ -325,17 +344,17 @@ class Diagram: f' "scale=trunc(iw/2)*2:trunc(ih/2)*2" -f mp4 "{path}"') # Remove files. - for i in frames: + for i in range(len(frames)): os.remove(self.sim.path / f"temp/img{i:05}.png") os.rmdir(self.sim.path / 'temp') print(f'Wrote to \"{path}\".', flush=True) -def render_frame_range(combo: Tuple[Diagram, List[int], str, int]) -> None: - self, frames, fol, num_frames = combo - for frame in frames: - self.generate_frame(frame, "save", fol) +def render_frame_range(combo: Tuple[Diagram, str, int, int, int]) -> None: + self, fol, offset, length, num_frames = combo + for i in range(length): + self.generate_frame(i, "save", fol, f"img{i+offset:05}.png") i = len(list((self.sim.path / fol).iterdir())) hashes = int(21*i/num_frames) print(f'Generating frames... |{"#"*hashes}{" "*(20-hashes)}|' + \ diff --git a/squish/simulation.py b/squish/simulation.py index 06d1cf8..9b4d945 100644 --- a/squish/simulation.py +++ b/squish/simulation.py @@ -21,12 +21,15 @@ class Simulation: __slots__ = ['domain', 'energy', 'path', 'frames'] - def __init__(self, domain: DomainParams, energy: Energy, name: Optional[str] = None) -> None: + def __init__(self, domain: DomainParams, energy: Energy, \ + name: Optional[str, int] = None) -> None: self.domain, self.energy = domain, energy self.frames = [] if name is None: self.path = generate_filepath(self, OUTPUT_DIR) + elif isinstance(name, int): + self.path = generate_filepath(self, OUTPUT_DIR, name) else: self.path = OUTPUT_DIR / name @@ -86,16 +89,16 @@ class Simulation: return distinct_count - def save(self, info: Dict) -> None: + def save(self, info: Dict, overwrite: bool = False) -> None: self.path.mkdir(exist_ok=True) path = self.path / 'data.squish' - with open(path, 'ab') as out: + with open(path, 'wb' if overwrite else 'ab') as out: pickle.dump(info, out, pickle.HIGHEST_PROTOCOL) def save_all(self) -> None: - self.save(self.initial_data) + self.save(self.initial_data, True) for i in range(len(self.frames)): self.save(self.frame_data(i)) @@ -175,7 +178,7 @@ class Flow(Simulation): def run(self, save: bool, log: bool, log_steps: int) -> None: if log: print(f"Find - {self.domain}", flush=True) - if save: self.save(self.initial_data) + if save: self.save(self.initial_data, True) if len(self) == 0: self.add_frame() i, grad_norm = 0, float('inf') @@ -212,7 +215,7 @@ class Flow(Simulation): self[i].add_sites(change/shrink_factor)) self.step_size /= shrink_factor - self.step_size = max(10e-4, self.step_size) + self.step_size = max(10e-6, self.step_size) self.frames.append(new_frame) @@ -269,7 +272,7 @@ class Search(Simulation): def run(self, save: bool, log: bool, log_steps: int) -> None: if log: print(f'Travel - {self.domain}', flush=True) - if save: self.save(self.initial_data) + if save: self.save(self.initial_data, True) if len(self) != 0: new_sites = self[0].site_arr @@ -350,7 +353,7 @@ class Shrink(Simulation): def run(self, save: bool, log: bool, log_steps: int) -> None: if log: print(f'Shrink - {self.domain}', flush=True) - if save: self.save(self.initial_data) + if save: self.save(self.initial_data, True) if len(self) != 0: new_sites = self[0].site_arr @@ -380,6 +383,4 @@ STR_TO_SIM = { "flow": Flow, "search": Search, "shrink": Shrink -} - -simulation = Simulation \ No newline at end of file +} \ No newline at end of file