Numpy lesson outline

numpy.py

import numpy as np

# Find local_stars.csv at:
# https://drive.google.com/file/d/1TKz0exDHL1d6a6qSU91QamjBgwL7EZ4_/view?usp=sharing
with open("local_stars.csv") as f:
    lines = f.read().split()

# Parse the header line.
columns = lines[0].split(",")
print(columns)

# Meaning of the columns:
# right_ascension: horizontal position along the celestial equator
# declination: vertical position above/below celestial equator
# parallax: amount the star's position in the sky changes throughout the year
# magnitude_g: how bright the star appears in our sky (inverse Logarithmic scale)
# magnitude_red: red-component of the color of the star
# magnitude_blue: blue-component of the color of the star

## Distance from our sun.

# Extract the parallax data for each star.
stars_parallax = list()

for line in lines[1:]:
    parallax = line.split(",")[2]
    stars_parallax.append(float(parallax))

# Convert to NumPy array.
stars_parallax = np.array(stars_parallax)

# Print the parallax for the first 10 stars. This is how slicing works.
stars_parallax[:10]  # in milli-arcseconds

# Do some basic math: compute distance.
stars_distance = 1 / stars_parallax * 1000  # in parsecs
stars_distance_ly = stars_distance * 3.26156  # in light years

# How far is the closest star?
print(stars_distance_ly.min())

## Loading the entire file.

# Read the entire file using numpy.
data = np.loadtxt("local_stars.csv", delimiter=",", skiprows=1)

# Extract all columns. This is how slicing works in two dimensions.
stars_ra = data[:, 0]
stars_dec = data[:, 1]
stars_parallax = data[:, 2]  # we already had this one, but for completeness
stars_mag = data[:, 3]
stars_red = data[:, 4]
stars_blue = data[:, 5]

# Question: Alkaid, Mizar form the tail end of the "big dipper".
#           How far are these stars from one another?
#                                                  * Dubhe
#                * Mizar         Megrez
#               /         *         *
#    * Alkaid  /       Alioth                          * Merak
#     \       /                          *
#      \     /                         Phecda
#       \   /
#        \ /
#         * Sol
#
# This is easier to answer if we transform the coordinates from
# Gaia-CRF (ra, dec, distance) -> Cartesian (x, y, z)
stars_x = stars_distance_ly * np.cos(stars_dec) * np.cos(stars_ra)
stars_y = stars_distance_ly * np.cos(stars_dec) * np.sin(stars_ra)
stars_z = stars_distance_ly * np.sin(stars_dec)

# Pack the three vectors together in a matrix (2D array).
# This is how we create bigger arrays from smaller ones.
stars_xyz = np.array([stars_x, stars_y, stars_z])

# Look at the "shape" of the array.
print(stars_xyz.shape)

# What if we want the stars along the rows and the XYZ along the columns?
# This is how transposing works.
stars_xyz = stars_xyz.T

# Now we can compute distance using Pythagoras.
# TODO: find the actual indices belonging to these stars.
distance_alkaid_mizar = np.sqrt(np.sum((stars_xyz[42, :] - stars_xyz[67, :]) ** 2))

# But there is a build-in function for it (there usually is).
distance_alkaid_mizar = np.linalg.norm(stars_xyz[42, :] - stars_xyz[67, :])


## Bonus: plot the stars in 3D (the data is begging for it!)
import pyvista as pv
from pyvistaqt import BackgroundPlotter

plotter = BackgroundPlotter()
plotter.set_background("black")
plotter.add_mesh(
    pv.PolyData(stars_xyz),
    render_points_as_spheres=True,
    point_size=5,
    scalars=stars_mag,
    opacity=[1.0, 0.1],
    cmap="magma_r",
    scalar_bar_args={
        "title": "Apparent magnitude",
        "color": "white",
    },
)
plotter.add_axes()

# Place the camera in the middle, looking along the celestial equator.
plotter.camera.position = (0, 0, 0)  # camera location
plotter.camera.focal_point = (0.01, 0, 0)  # look along +x
plotter.camera.up = (0, 0, 1)  # z-axis is up

# Add the Sun. This sphere is initially hidden because the camera is inside of it, so
# you can explore the night sky as it appears on Earth. If you zoom out far enough, the
# Sun-sphere becomes visible to provide a point of reference when exploring our stellar
# neighbourhood.
sun = pv.Sphere(radius=0.2, center=(0, 0, 0))
plotter.add_mesh(sun, color="white")


## Making a Hertzsprung-Russell diagram: color versus magnitude.

# Compute magnitude as if the star was at 10 parsecs from us.
# Formula for this is taken from scientific papers.
stars_mag_norm = stars_mag - 5 * np.log10(stars_distance) + 5

# Compute where the color of the star falls in the range from blue to red.
stars_color = stars_blue - stars_red

# Make the HRD plot.
import matplotlib.pyplot as plt

plt.figure(figsize=(6, 8))
plt.scatter(
    stars_color, stars_mag_norm, s=1, c=stars_distance_ly, cmap="plasma", alpha=0.1
)
cb = plt.colorbar(label="Distance (light years)", fraction=0.05, aspect=40)
cb.solids.set(alpha=1)  # make the colors in the colorbar not be transparent
plt.gca().invert_yaxis()
plt.title("Hertzsprung-Russell diagram")
plt.xlabel("Color (G_bp - G_rp)")
plt.ylabel("Magnitude (M_g)")
plt.tight_layout()

# Extract the main sequence curve. Now we get into some more advanced NumPy.
bins = np.linspace(0.5, 4, 100)  # create equally spaced numbers with `linspace`
idx = np.digitize(stars_color, bins)  # assign stars to these bins
median_m_per_bin = np.array(
    [np.median(stars_mag_norm[idx == i]) for i in range(1, len(bins))]
)  # this is how you select things by boolean mask
coeffs = np.polyfit(bins[:-1], median_m_per_bin, deg=5)  # 5th order polynomial

# Evaluate the polynomial fit and compare against the extracted main sequence.
plt.plot(bins[:-1], median_m_per_bin, color="black", label="main sequence curve")
plt.plot(bins, np.polyval(coeffs, bins), color="yellow", label="polynomial fit")


def estimate_dist(apparent_mag, color):
    """Estimate the distance to a star given its magnitude and color. In parsecs."""
    est_true_mag = np.polyval(coeffs, color)
    return np.pow(10, (apparent_mag - est_true_mag) / 5 + 1)


# Load some far away stars of the "main sequence" type.
# Find far_stars.csv at:
# https://drive.google.com/file/d/1vB281gaPzJgs-jeqkPCSm3ypOaGjMnTs/view?usp=sharing
data = np.loadtxt("far_stars.csv", delimiter=",", skiprows=1)
# Bonus: this is how list unpacking works. Rows are assigned to variables.
far_parallax, far_mag, far_red, far_blue = data[:, 2:6].T

# Estimate the distance of these stars based on their apparent magnitude and color.
far_distance = estimate_dist(far_mag, far_blue - far_red)

# Plot parallax-estimated versus color-estimated distance.
plt.figure()
plt.scatter(1000 / far_parallax, far_distance)
plt.plot(np.linspace(0, 7000, 10), np.linspace(0, 7000, 10), color="black")
plt.xlabel("Distance from parallax (parsecs)")
plt.ylabel("Distance from main-sequence fitting (parsecs)")
plt.gca().set_aspect("equal")