sims_python_files/visualisation_swift_monofonic_tests.py at main · glatterf42/sims_python_files · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# import vtk
import sys
import time
from pathlib import Path
from sys import getsizeof

import h5py
import numba
import numpy as np
import scipy.spatial


@numba.njit()
def determine_desired_range(offset, minimum, upper_limit_bottom, lower_limit_top, maximum):
    a = minimum
    b = maximum

    if offset < 0:
        a = lower_limit_top
    elif offset > 0:
        b = upper_limit_bottom

    return a, b


@numba.njit()
def find_coordinates_to_move(minimum, maximum, x_offset, y_offset, z_offset, move_candidates):
    coordinates_to_move = []
    x_start, x_end = determine_desired_range(x_offset, minimum, upper_limit_bottom, lower_limit_top, maximum)
    y_start, y_end = determine_desired_range(y_offset, minimum, upper_limit_bottom, lower_limit_top, maximum)
    z_start, z_end = determine_desired_range(z_offset, minimum, upper_limit_bottom, lower_limit_top, maximum)

    for particle in move_candidates:
        point = particle[0:3]
        if x_start <= point[0] <= x_end and y_start <= point[1] <= y_end and z_start <= point[2] <= z_end:
            coordinates_to_move.append(particle)

    return coordinates_to_move


directory = Path(r"/home/ben/sims/data_swift/monofonic_tests/shannon_256_100/")

for filename in sorted(directory.glob("output_0004.hdf5")):
    print(filename)
    file = h5py.File(str(filename), "r")
    Header = file['Header']

    original_coordinates = file["PartType1"]["Coordinates"][:]  # for cdm particles
    names = file["PartType1"]["ParticleIDs"][:]
    velocities = file["PartType1"]["Velocities"][:]
    masses = file['PartType1']['Masses'][:]
    group_ids = file['PartType1']['FOFGroupIDs'][:]
    absolute_velo = np.sqrt(np.sum(velocities ** 2, axis=1))

    original_data = np.vstack([
        original_coordinates[::, 0],
        original_coordinates[::, 1],
        original_coordinates[::, 2],
        names,
        velocities[::, 0],
        velocities[::, 1],
        velocities[::, 2],
        masses,
        group_ids,
        absolute_velo,
    ]).T
    print(original_data.shape)
    assert (original_coordinates == original_data[::, 0:3]).all()

    boundaries = Header.attrs['BoxSize']  # BoxLength for e5 boxes depends on Nres, 2.36438 for 256, 4.72876 for 512.
    print(boundaries, len(names))
    if not boundaries.shape:
        boundaries = np.array([boundaries] * 3)
    offsets = [-1, 0, 1]
    transformed_data = original_data[:]
    number_of_time_that_points_have_been_found = 0

    # assumes cube form and 0.1 as desired ratio to move
    minimum = 0.0
    maximum = max(boundaries)
    box_length = maximum - minimum
    # magic number: mean particle separation * 15 in units of box_length
    range_to_move = 0.0586 * box_length  # mean particle separation: 0.78125 Mpc for 128, 0.390625 Mpc for 256, etc
    upper_limit_bottom = minimum + range_to_move
    lower_limit_top = maximum - range_to_move

    print("Find candidates to move...")


    @numba.njit()
    def find_move_candidates():
        move_candidates = []
        print("finding move candidates")
        for particle in original_data:
            point = particle[0:3]
            if (
                    minimum <= point[0] <= upper_limit_bottom or
                    lower_limit_top <= point[0] <= maximum or
                    minimum <= point[1] <= upper_limit_bottom or
                    lower_limit_top <= point[1] <= maximum or
                    minimum <= point[2] <= upper_limit_bottom or
                    lower_limit_top <= point[2] <= maximum
            ):
                move_candidates.append(particle)
                # print(point)
        return move_candidates


    move_candidates = find_move_candidates()
    move_candidates = np.array(move_candidates)

    print("...done.")
    for x in offsets:
        for y in offsets:
            for z in offsets:
                if (x, y, z) == (0, 0, 0):
                    continue
                moved_coordinates = find_coordinates_to_move(minimum, maximum, x, y, z, move_candidates)
                # print(moved_coordinates)
                moved_coordinates = np.array(moved_coordinates)
                # if not moved_coordinates.all():
                #     print(f"nothing moved in {(x,y,z)}")
                #     continue
                moved_coordinates[::, 0] += x * boundaries[0]
                moved_coordinates[::, 1] += y * boundaries[1]
                moved_coordinates[::, 2] += z * boundaries[2]
                transformed_data = np.vstack((transformed_data, moved_coordinates))
                number_of_time_that_points_have_been_found += 1
                print(f"Points found: {number_of_time_that_points_have_been_found}/26...")

    # assert coordinates.shape[0] == original_coordinates.shape[0] * 3 ** 3 #check that the new space has the shape we want it to have

    num_nearest_neighbors = 40
    print("Building 3d-Tree for all particles...")
    coordinates = transformed_data[::, 0:3]
    print(coordinates.shape)
    tree = scipy.spatial.KDTree(coordinates)
    print(getsizeof(tree) / 1024, "KB")
    print("...done.")
    print("Searching neighbours...")
    a = time.perf_counter_ns()
    distances, indices = tree.query([coordinates], k=num_nearest_neighbors, workers=6)
    # shape of closest_neighbours: (1, xxxx, 40)
    b = time.perf_counter_ns()
    print("...found neighbours.")
    print(f"took {(b - a) / 1000 / 1000:.2f} ms")
    distances = distances[0]  # to (xxxx, 40)
    indices = indices[0]  # to (xxxx, 40)
    print(distances.shape)
    print(indices.shape)
    print(indices)
    mass_array = []
    print("fetching masses")
    for subindices in indices:  # subindices is (40)
        # can maybe be optimized to remove loop
        masses = transformed_data[subindices, 7]
        mass_array.append(masses)
    mass_array = np.array(mass_array)
    print("finished fetching masses")
    # print(closest_neighbours, indices)
    # print(indices)

    # densities = num_nearest_neighbors * mass_per_particle / np.mean(closest_neighbours, axis=1) ** 3
    total_masses = np.sum(mass_array, axis=1)
    densities = total_masses / np.mean(distances, axis=1) ** 3
    alt_densities = total_masses / np.max(distances, axis=1) ** 3

    # print(closest_neighbours.shape)

    # print(densities)
    # print(densities.shape)
    all_data = np.column_stack([list(range(densities.shape[0])), transformed_data, densities, alt_densities])
    # print(all_data.shape)
    # print(original_data.shape[0])
    export_data = all_data[:original_data.shape[0]]
    # print(export_data.shape)

    # all_data = np.append(coordinates, velocities, axis=1)
    # all_data = np.column_stack((all_data, absolute_velo, names, densities))
    # sorted_index = np.argsort(all_data[::, 7], kind="stable")
    # all_data = all_data[sorted_index, :]

    #    np.savetxt("out_"+filename.with_suffix(".csv").name, all_data[indices], delimiter=",", fmt="%.3f", header="x,y,z,vx,vy,vz,v,name") #if indices are needed
    np.savetxt(directory / f"visualisation_{filename.with_suffix('.csv').name}",
               export_data,
               delimiter=",",
               fmt="%.3f",
               header="num,x,y,z,name,vx,vy,vz,masse,groupid,v,density,density_alt")
    file.close()