Mass to Drag by Density and Radii (Comparative Calculator)

if samp_spheres == 0:
    spec_spheres = int(input('Number of exact spheres:'))
    if spec_spheres == 1:
        radii = input(f'Radius of sphere in micrometers:')
        radii = [float(radii) if radii else 69]
    elif spec_spheres > 1:
        radii = [float(input(f'Radius of sphere {i+1} in micrometers:')) for i in range(spec_spheres)]  # manual
elif samp_spheres > 0:
    np.random.seed(samp_spheres)
    radii = uniform.rvs(1, 1e5, size=samp_spheres)  # automatic

from math import pi
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import uniform

def main():
    '''This function takes these inputs:
    
        Enter on new lines:

        1. density of sphere,
        2. density of fluid,
        3. velocity of motion of the sphere relative to the fluid,
        4. a number of example spheres or 0 for exact spheres,
        5. a number of exact spheres (if 0 on 4),
        6. length of radius or radii of exact spheres (if 0 on 4 and not 0 on 5,).

        It plots Mass to Drag ratio and Kinetic Energy relative to radius length.'''
    
    sphere_density = input('Sphere Density in kg/m^3 (water ice is around 920):')
    sphere_density = float(sphere_density) if sphere_density else 920
    fluid_density = input('Fluid Density in kg/m^3 (the atmosphere at sea level is around 1.21):')
    fluid_density = float(fluid_density) if fluid_density else 1.21
    velocity = input('Velocity in m/s (gravity acceleration at sea level is around 9.8 m/s^2):')
    velocity = float(velocity) if velocity else .8
    samp_spheres = input('Number of example spheres (0 for exact):')
    samp_spheres = int(samp_spheres) if samp_spheres else 0
    
    if samp_spheres == 0:
        spec_spheres = input('Number of exact spheres:')
        spec_spheres = int(spec_spheres) if spec_spheres else 5
        if spec_spheres == 1:
            radii = input(f'Radius of sphere in micrometers:')
            radii = [float(radii) if radii else 69]
        elif spec_spheres > 1:
            radii = [input(f'Radius of sphere {i+1} in micrometers:') for i in range(spec_spheres)]
            radii = [float(r) if r else np.random.uniform(1, 1e5) for r in radii]
    elif samp_spheres > 0:
        np.random.seed(samp_spheres)
        radii = uniform.rvs(1, 1e5, size=samp_spheres)  # automatic
    
    radii.sort()
    radii_normz = norm(radii)
    vols = sphere_volume(radii)  # um^3
    masses = mass(sphere_density, vols)  # kg
    kinergy = [(0.5 * mass * velocity**2) for mass in masses]
    ke_normz = norm(kinergy)
    ke_mean = np.mean(kinergy)
    drags = sphere_drag(fluid_density, velocity, radii)  # m*s^3
    mtds = mass_to_drag(masses, drags)  # kg/m^2
    mtds_normz = norm(mtds)
    mini, maxi = min(radii_normz), max(radii_normz)
    mtd_mean = np.mean(mtds)
    alpha = 0.5
    
    # Create a pandas DataFrame
    data = {'Radius': radii, 'Mass_to_Drag_Ratio': mtds, 'Kinergy': kinergy, 'Normed_Radius': radii_normz, 'Normed_Kinergies': ke_normz}
    df = pd.DataFrame(data)
    corcof = df['Radius'].corr(df['Mass_to_Drag_Ratio'])
    print(f'Correlation Coefficient: {corcof}')
    print(f'Average Mass to Drag Ratio (exact if 1 sphere): {mtd_mean:.3g} kg/m*s^3')
    print(f'Average kinetic energy (exact if 1 sphere): {ke_mean:.3g} Joules')


    # Create subplots with shared x-axis
    fig, ax1 = plt.subplots(figsize=(12, 9))

    # Scatter plot for Mass to Drag Ratio
    sns.scatterplot(data=df, x='Radius', y='Mass_to_Drag_Ratio', hue='Normed_Radius', size=radii_normz, sizes=(mini, maxi), alpha=alpha, palette='blend:#7AB,#EDB', legend=False, ax=ax1)

    # Create a second y-axis for Kinergy
    ax2 = ax1.twinx()
    ax2.scatter(df['Radius'], df['Kinergy'], marker='o', color='black', alpha=alpha, label='Kinetic Energy')

    # Set y-axis limits and label for the first axis
    ax1.set_ylim(0, 1.5 * max(mtds_normz))
    ax1.set_ylabel('Mass to Drag Ratio (kg/m*s^3)', fontsize=14)

    # Set y-axis limits and label for the second axis
    ax2.set_ylim(0, 1.1 * max(ke_normz))
    ax2.set_ylabel('Kinetic Energy in Joules (kg*(m/s)^2)', fontsize=14)

    # Set title
    plt.title('Mass to Drag Ratio and Kinetic Energy vs. Sphere Radius', fontsize=16)

    # Set x-axis label
    ax1.set_xlabel('Radius (micrometers)', fontsize=14)

    # Display legend
    ax2.legend(loc='upper right')

    plt.xlim(0, 1.5 * max(radii_normz))
    plt.show()


def norm(radii):
    '''Returns normalized radii for proportionality in sns.scatterplot size param.'''
    normz = [radius / max(radii) for radius in radii]    
    radii_normz = np.array([radius * norm for radius, norm in zip(radii, normz)])
    return radii_normz

def sphere_volume(radii):
    '''Returns sphere volume after unit conversion um^3 to m^3'''
    return [(4/3*pi*radius**3 * 1e-18) for radius in radii]

def sphere_drag(fluid_density, velocity, radii):
    '''Returns sphere drag after unit conversion um^2 to m^2'''
    return [(0.45*0.5*fluid_density*(velocity**2)*pi*radius**2 * 1e-12) for radius in radii]
# D = Cd * . 5 * rho * V^2 * A. Where D is equal to the drag (m/s^2), rho is the fluid density (kg/m^3), V is the velocity (m/s), A is a reference area (m^2), and Cd is the drag coefficient (.45 for a sphere).
def mass(sphere_density, vols):
    '''Returns mass as kg'''
    return [(sphere_density * vol) for vol in vols]

def mass_to_drag(masses, drags):
    '''Returns mass to drag ratio as kg/m*s^3'''
    return [mass/drag for mass, drag in zip(masses, drags)]

if __name__ == '__main__':
    main()