Refactor to remove extra gc.collect()

This commit is contained in:
LilyRose2798 2024-04-15 22:59:29 +10:00
parent 27a8d2584d
commit 3c6063979c
1 changed files with 50 additions and 50 deletions

View File

@ -15,6 +15,7 @@ from typing import Literal, TypeVar
from png import Writer
from cmap import Colormap
from hilbert import decode
from numpy.typing import NDArray
import numpy as np
ip_bytes = 4
@ -157,27 +158,27 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
Writer(tile_size, tile_size, greyscale = False, alpha = alpha).write_packed(path.open("wb"), rows)
print("done")
print(f"reading scan data from file '{input_path}'...", end = " ", flush = True)
data = np.fromfile(input_path, dtype = np.uint32).reshape(-1, 2)
ip_arr = np.copy(data.T[0])
rtt_arr = np.copy(data.T[1])
print("done")
del data
collect()
def get_scan_data() -> tuple[NDArray[np.uint32], NDArray[np.uint32]]:
print(f"reading scan data from file '{input_path}'...", end = " ", flush = True)
data = np.fromfile(input_path, dtype = np.uint32).reshape(-1, 2)
ip_arr = np.copy(data.T[0])
rtt_arr = np.copy(data.T[1])
print("done")
return (ip_arr, rtt_arr)
print(f"reading coordinates from file '{coords_path}'...", end = " ", flush = True)
ip_coords = np.fromfile(coords_path, dtype = np.uint16).reshape(-1, 2)
print("done")
print(f"converting ip addresses to coordinates...", end = " ", flush = True)
xs, ys = ip_coords[ip_arr].T
coords = (ys, xs)
print("done")
del ip_coords
del ip_arr
collect()
def get_all_data() -> tuple[tuple[NDArray[np.uint16], NDArray[np.uint16]], NDArray[np.uint32]]:
ip_arr, rtt_arr = get_scan_data()
print(f"reading coordinates from file '{coords_path}'...", end = " ", flush = True)
ip_coords = np.fromfile(coords_path, dtype = np.uint16).reshape(-1, 2)
print("done")
print(f"converting ip addresses to coordinates...", end = " ", flush = True)
xs, ys = ip_coords[ip_arr].T
print("done")
return ((ys, xs), rtt_arr)
coords, rtt_arr = get_all_data()
def generate_density():
variant_name = "density"
possible_overlaps = 1
print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array of density data...", end = " ", flush = True)
@ -205,7 +206,7 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
def write_all_colormaps():
for colormap_name, colormap in colormaps:
create_images(density_data, colormap, possible_overlaps, tiles_dir / variant_name / colormap_name)
create_images(density_data, colormap, possible_overlaps, tiles_dir / "density" / colormap_name)
write_all_colormaps()
while density_data.shape[0] > tile_size:
@ -213,31 +214,32 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
write_all_colormaps()
def generate_rtt():
nonlocal rtt_arr
variant_name = "rtt"
num_colors = (1 << 16) - 1
multiplier = num_colors - 1
print(f"retrieving {quantile:.1%} quantile for rtt data...", end = " ", flush = True)
rtt_quantile = np.quantile(rtt_arr, quantile)
print("done")
print(f"scaling rtt data using rtt quantile...", end = " ", flush = True)
rtt_arr_f = rtt_arr / rtt_quantile
print("done")
del rtt_arr
collect()
print("clipping rtt data between 0 and 1...", end = " ", flush = True)
rtt_arr_f.clip(0, 1, out = rtt_arr_f)
print("done")
def get_rtt_data():
nonlocal rtt_arr
print(f"retrieving {quantile:.1%} quantile for rtt data...", end = " ", flush = True)
rtt_quantile = np.quantile(rtt_arr, quantile)
print("done")
print(f"scaling rtt data using rtt quantile...", end = " ", flush = True)
rtt_arr_f = rtt_arr / rtt_quantile
print("done")
del rtt_arr
collect()
print("clipping rtt data between 0 and 1...", end = " ", flush = True)
rtt_arr_f.clip(0, 1, out = rtt_arr_f)
print("done")
print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array for rtt data...", end = " ", flush = True)
rtt_data = np.full((num_ips_sqrt, num_ips_sqrt), np.nan, dtype = np.float32)
print("done")
print(f"assigning values to rtt data array...", end = " ", flush = True)
rtt_data[coords] = rtt_arr_f
print("done")
del rtt_arr_f
collect()
print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array for rtt data...", end = " ", flush = True)
rtt_data = np.full((num_ips_sqrt, num_ips_sqrt), np.nan, dtype = np.float32)
print("done")
print(f"assigning values to rtt data array...", end = " ", flush = True)
rtt_data[coords] = rtt_arr_f
print("done")
return rtt_data
rtt_data = get_rtt_data()
def squish():
nonlocal rtt_data
@ -278,7 +280,7 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
for _ in range(skip_iters):
squish()
def write_all_colormaps():
def get_normalized_data():
print(f"normalizing rtt data: multiplying...", end = " ", flush = True)
rtt_data_f = rtt_data * multiplier
print(f"incrementing...", end = " ", flush = True)
@ -287,12 +289,14 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
# rtt_data_f[np.isnan(rtt_data_f)] = 0.0
print(f"converting to ints...", end = " ", flush = True)
with catch_warnings():
rtt_data_i = rtt_data_f.astype(np.uint16)
rtt_data_norm = rtt_data_f.astype(np.uint16)
print("done")
del rtt_data_f
collect()
return rtt_data_norm
def write_all_colormaps():
rtt_data_norm = get_normalized_data()
for colormap_name, colormap in colormaps:
create_images(rtt_data_i, colormap, num_colors, tiles_dir / variant_name / colormap_name)
create_images(rtt_data_norm, colormap, num_colors, tiles_dir / "rtt" / colormap_name)
write_all_colormaps()
while rtt_data.shape[0] > tile_size:
@ -306,10 +310,6 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
collect()
if should_generate_density:
generate_density()
del xs
del ys
del coords
collect()
if json_path is not None and tiles_dir_parts is not None:
try: