diff --git a/remap_particle_IDs.py b/remap_particle_IDs.py index 8570272..32efed9 100644 --- a/remap_particle_IDs.py +++ b/remap_particle_IDs.py @@ -1,27 +1,48 @@ +from itertools import product + import numpy as np -def convert_ID_between_resolutions(particle_ID: int, Nres_from: int, Nres_to: int): - particle_k = particle_ID % Nres_from - particle_j = ((particle_ID - particle_k) // Nres_from) % Nres_from - particle_i = ((particle_ID - particle_k) // Nres_from - particle_j) // Nres_from - lowres_ID = ((particle_i * Nres_to) + particle_j) * Nres_to + particle_k - return lowres_ID +def original_position(Nres: int, particle_ID: int): + particle_k = particle_ID % Nres + particle_j = ((particle_ID - particle_k) // Nres) % Nres + particle_i = ((particle_ID - particle_k) // Nres - particle_j) // Nres + return np.array([particle_i, particle_j, particle_k]) + + +def upscale_IDs(particle_ID: int, Nres_min: int, Nres_max: int): + assert Nres_max % Nres_min == 0 + N = Nres_max // Nres_min + + orig = original_position(Nres_min, particle_ID) + mult = orig * N + for shift in product(range(N), repeat=3): + variant = mult + np.array(shift) + yield ((variant[0] * Nres_max) + variant[1]) * Nres_max + variant[2] + + +def downscale_IDs(particle_ID: int, Nres_max: int, Nres_min: int): + assert Nres_max % Nres_min == 0 + N = Nres_max // Nres_min + + orig = original_position(Nres_max, particle_ID) + mult = np.floor_divide(orig, N) + return ((mult[0] * Nres_min) + mult[1]) * Nres_min + mult[2] if __name__ == "__main__": - test_particle = np.array([0, 127, 127]) + test_particle = np.array([0, 0, 127]) # maximum_test = np.array([127, 127, 127]) #this works, Nres - 1 is the maximum for (i,j,k) Nres_1 = 128 Nres_2 = 256 - particle_ID_1 = ((test_particle[0] * Nres_1) + test_particle[1]) * Nres_1 + test_particle[2] - particle_ID_2 = ((test_particle[0] * Nres_2) + test_particle[1]) * Nres_2 + test_particle[2] + test_particle_id = ((test_particle[0] * Nres_1) + test_particle[1]) * Nres_1 + test_particle[2] + print(test_particle_id) - print(particle_ID_1, particle_ID_2) + particle_ID_1_converted = upscale_IDs(test_particle_id, Nres_1, Nres_2) - particle_ID_1_converted = convert_ID_between_resolutions(particle_ID_1, Nres_1, Nres_2) - particle_ID_1_converted_converted = convert_ID_between_resolutions(particle_ID_1_converted, Nres_2, Nres_1) - - print(particle_ID_1, particle_ID_1_converted, particle_ID_1_converted_converted) + for id in particle_ID_1_converted: + reverse = downscale_IDs(id, Nres_2, Nres_1) + print(id, reverse) + assert reverse == test_particle_id