Changed diagram.py so that parallelization doesn't use excess memory, and other small tweaks

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.

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)

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)}|' + \

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
@@ -381,5 +384,3 @@ STR_TO_SIM = {
 	"search": Search,
 	"shrink": Shrink
 }
-
-simulation = Simulation