LMC-Structural-Parameter-Estimation/utils.py at main · Deepnika/LMC-Structural-Parameter-Estimation · 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
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from astropy.coordinates import SkyCoord

# Model
def create_disk(N, scale_length, scale_height, ell_factor):
    r = np.random.gamma(shape=2, scale=scale_length, size = N)
    theta = np.random.uniform(low=0, high=360, size = N)

    xi = ell_factor * r * np.cos(np.radians(theta))
    eta = r * np.sin(np.radians(theta))
    zeta = np.random.laplace(loc=0, scale=scale_height, size = N)

    x = r * np.cos(np.radians(theta))
    y = r * np.sin(np.radians(theta))
    return xi, eta, zeta, x, y, zeta

def first_rotation(xi, eta, zeta, theta_ma):
    R_rot1 = np.array([[np.cos(np.radians(theta_ma)), -np.sin(np.radians(theta_ma)), 0],
                    [np.sin(np.radians(theta_ma)), np.cos(np.radians(theta_ma)), 0],
                    [0, 0, 1]])
    vector = np.stack([xi,eta,zeta],axis=1)
    result = np.dot(R_rot1,vector.T).T
    return result[:,0], result[:,1], result[:,2]

def second_rotation(xi_rot1, eta_rot1, zeta_rot1, theta_lon, i_incl):
    R_rot2 = np.array([[np.cos(np.radians(theta_lon)), -np.sin(np.radians(theta_lon))*np.cos(np.radians(i_incl)), -np.sin(np.radians(theta_lon))*np.sin(np.radians(i_incl))],
                    [np.sin(np.radians(theta_lon)), np.cos(np.radians(theta_lon))*np.cos(np.radians(i_incl)), np.cos(np.radians(theta_lon))*np.sin(np.radians(i_incl))],
                    [0, -np.sin(np.radians(i_incl)), np.cos(np.radians(i_incl))]])
    vector = np.stack([xi_rot1,eta_rot1,zeta_rot1],axis=1)
    result = np.dot(R_rot2,vector.T).T
    return result[:,0], result[:,1], result[:,2]

def rect_heliocentric_frame(xi_rot2, eta_rot2, zeta_rot2, dist_LMC, alpha_c_LMC, delta_c_LMC):
    R_proj = np.array([[np.sin(np.radians(alpha_c_LMC)), -np.cos(np.radians(alpha_c_LMC))*np.sin(np.radians(delta_c_LMC)), -np.cos(np.radians(alpha_c_LMC))*np.cos(np.radians(delta_c_LMC))],
                       [-np.cos(np.radians(alpha_c_LMC)), -np.sin(np.radians(alpha_c_LMC))*np.sin(np.radians(delta_c_LMC)), -np.sin(np.radians(alpha_c_LMC))*np.cos(np.radians(delta_c_LMC))],
                       [0, np.cos(np.radians(delta_c_LMC)), -np.sin(np.radians(delta_c_LMC))]])
    T_proj = np.array([dist_LMC*np.cos(np.radians(delta_c_LMC))*np.cos(np.radians(alpha_c_LMC)),
                       dist_LMC*np.cos(np.radians(delta_c_LMC))*np.sin(np.radians(alpha_c_LMC)),
                       dist_LMC*np.sin(np.radians(delta_c_LMC))])
    vector = np.stack([xi_rot2,eta_rot2,zeta_rot2],axis=1)
    result = np.dot(R_proj,vector.T).T + T_proj
    return result[:,0], result[:,1], result[:,2]

def gaia_observables(x,y,z):
    parallax = 1/np.sqrt(x**2+y**2+z**2)
    ra = np.degrees(np.arctan2(y,x))
    dec = np.degrees(np.arcsin(z/(1/parallax)))
    return parallax, ra, dec

def plot_grouped_histograms(data, group_column, bins=50):
    unique_groups = data[group_column].unique()
    numeric_columns = data.select_dtypes(include=['float64', 'int64']).columns
    num_columns = len(numeric_columns)

    rows = int(np.ceil(num_columns / 3))  # 3 columns per row
    fig, axes = plt.subplots(rows, 3, figsize=(15, 5 * rows))
    axes = axes.flatten()

    for i, column in enumerate(numeric_columns):
        for group in unique_groups:
            group_data = data[data[group_column] == group]
            axes[i].hist(group_data[column], bins=bins, alpha=0.5, label=str(group), density=True)

        axes[i].set_title(f"{column}")
        axes[i].set_xlabel(column)
        axes[i].set_ylabel("Frequency")
        axes[i].legend(title=group_column)
        axes[i].grid(True)

    # Hide any unused subplots
    for j in range(i + 1, len(axes)):
        fig.delaxes(axes[j])

    plt.tight_layout()
    plt.show()

def get_cartesian_coordinates(df):
    ra = df["ra"].values  # degrees
    dec = df["dec"].values  # degrees
    parallax = df["parallax"].values  # milliarcseconds (mas)

    # Convert parallax to distance in kpc
    # 1/parallax [arcsec] = distance [pc] => 1000/parallax [mas] = distance [pc] => /1000 to get kpc
    distance_kpc = 1 / parallax
    coords = SkyCoord(ra=ra*u.deg, dec=dec*u.deg, distance=distance_kpc*u.kpc)
    x = coords.cartesian.x.value
    y = coords.cartesian.y.value
    z = coords.cartesian.z.value

    return x, y, z

def plot_3d_distribution(df):
    fig = go.Figure(data=[go.Scatter3d(
        x = df["x"], y = df["y"], z = df["z"],
        mode='markers',
        marker=dict(
            size=1,
            color=df["parallax"],
            colorscale='Viridis',
            colorbar=dict(title='Parallax (mas)'),
            opacity=0.7
        )
    )])
    fig.update_layout(
        scene=dict(
            xaxis_title='X (kpc)',
            yaxis_title='Y (kpc)',
            zaxis_title='Z (kpc)'
        ),
        width=1000,
        height=800,
        title='3D Distribution of Simulated LMC Stars'
    )
    return fig