Skip to content

fix main_density_calc to reduce memory usage #47

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shruthivis merged 1 commit into from
Apr 13, 2026
Merged

fix main_density_calc to reduce memory usage #47

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 11 additions & 10 deletions src/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import pandas as pd
import os
from multiprocessing import Pool
from functools import partial
from sparse_grid import SparseGrid
from bead_density import BeadDensity
from patch_computer import calc_bead_spread, get_patches, annotate_patches
Expand All @@ -11,13 +10,13 @@
import argparse
import tqdm

def main_density_calc(i, coords, mass, radius, grid, voxel_size, n_breaks):
bead_density = BeadDensity(coords.shape[0], grid=grid, voxel_size=voxel_size)
# Obtain min-max coords for each bead across all models to construct a kernel.
def main_density_calc(coords, mass, radius, bead_density, n_breaks):
# bead_density = BeadDensity(coords.shape[0], grid=grid, voxel_size=voxel_size)
# Obtain min-max coords for the bead across all models to construct a kernel.
# k1 --> min xyz coords of kernel; k2 --> max xyz coords of kernel.
k1, k2 = _get_bounding_box(coords[:,i,:])
k1, k2 = _get_bounding_box(coords)
bead_density.construct_kernel(k1,k2)
return bead_density.return_density_opt(coords[:,i,:], radius[i], mass[i], n_breaks)
return bead_density.return_density_opt(coords, radius, mass, n_breaks)

def scale(v):
return (v - min(v)) / (max(v) - min(v))
Expand All @@ -37,9 +36,10 @@ def get_file_type(input):
i = i+1
return file_type

def get_bead_spread(i, coords, mass, radius, grid, voxel_size, n_breaks):
def get_bead_spread(arguments):
coords, mass, radius, grid, bead_density, n_breaks = arguments
# dummy function to call main_density_calc for parallelization
density = main_density_calc(i, coords, mass, radius, grid, voxel_size, n_breaks)
density = main_density_calc(coords, mass, radius, bead_density, n_breaks)
spread = calc_bead_spread(density, grid)
return spread

Expand All @@ -65,10 +65,11 @@ def run_prism( coords, mass, radius, ps_names, args, output_dir = None ):
chunksize, extra = divmod(coords.shape[1], cores_ * 4)
if extra:
chunksize += 1

bead_density = BeadDensity(coords.shape[0], grid=grid, voxel_size=args.voxel_size)
with Pool(cores_) as p:
bead_spread = []
for spread in tqdm.tqdm( p.imap( partial(get_bead_spread, coords=coords, mass=mass, radius=radius, grid=grid, voxel_size=args.voxel_size, n_breaks=args.n_breaks), range(coords.shape[1]), chunksize=chunksize ) ):
arguments = [(coords[:,i,:], mass[i], radius[i], grid, bead_density, args.n_breaks) for i in range(coords.shape[1])]
for spread in tqdm.tqdm( p.imap( get_bead_spread, arguments, chunksize=chunksize ) ):
bead_spread.append( spread )
bead_spread = scale(bead_spread)
print('Bead Spread calculation done')
Expand Down
Loading