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

100
ipmap.py
View File

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