from bs4 import BeautifulSoup import glob import math
import os
import re import requests
import tarfile
URL = "https://people.sissa.it/~sbressan/CAF09_V1.2S_M36_LT/"
TARGET_DIR = "src/access/"
MOD_TEMPLATE = """
// This code is generated by generate_code.py, do not modify it manually.
pub(crate) const PARSEC_URL: &str = "{URL}";
pub(crate) mod data;
pub(crate) mod metallicity;
pub(crate) mod masses;
"""
DATA_TEMPLATE = """
// This code is generated by generate_code.py, do not modify it manually.
use lazy_static::lazy_static;
use crate::data::ParsecData;
lazy_static! {{
{static_data}
}}
lazy_static! {{
pub(crate) static ref DATA: [&'static ParsecData; {array_size}] = [
{access_array}
];
}}
"""
METALLICITY_TEMPLATE = """
// This code is generated by generate_code.py, do not modify it manually.
pub(crate) static METALLICITIES_IN_MASS_FRACTION: [f64; {number_of_metallicities}] = [
{mass_fraction_array}
];
pub(crate) static METALLICITIES_IN_DEX: [f64; {number_of_metallicities}] = [
{dex_array}
];
pub(crate) static METALLICITY_NAMES: [&str; {number_of_metallicities}] = [
{names_array}
];
/// Using this, the crate knows which file to download during intialisation.
pub(crate) static METALLICITY_ARCHIVES: [&str; {number_of_metallicities}] = [
{archives_array}
];
"""
MASSES_TEMPLATE = """
// This code is generated by generate_code.py, do not modify it manually.
{sorted_masses}
{filenames}
pub(crate) static MASSES: [&[f64]; {number_of_metallicities}] = [
{metallicity_to_masses}
];
pub(crate) static FILENAMES: [&[&str]; {number_of_metallicities}] = [
{metallicity_to_filenames}
];
"""
def assure_dev_data_folder():
if not os.path.exists('dev_data'):
os.makedirs('dev_data')
def collect_archive_names():
FILE_PATH = "dev_data/archive_names.html"
if not os.path.isfile(FILE_PATH):
web_page = requests.get(URL).text
with open(FILE_PATH, 'w') as file:
file.write(web_page)
with open(FILE_PATH, 'r') as file:
web_page = file.read()
soup = BeautifulSoup(web_page, 'html.parser')
archive_names = []
for link in soup.find_all('a'):
href = link.get('href')
if href and href.endswith('.tar.gz'):
archive_names.append(href)
return archive_names
def assure_archive_downloaded(archive_name):
archive_path = f'dev_data/{archive_name}'
if not os.path.exists(archive_path):
archive_url = f'{URL}{archive_name}'
print(f'Downloading {archive_url}...')
with requests.get(archive_url, stream=True) as r:
with open(archive_path, 'wb') as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
def assure_extracted(archive_name):
archive_path = f'dev_data/{archive_name}'
extract_dir = 'dev_data'
target_path = extract_dir + '/' + archive_name.replace('.tar.gz', '')
if not os.path.exists(target_path):
with tarfile.open(archive_path, 'r:gz') as tar:
print(f'Extracting {archive_path}...')
tar.extractall(path=extract_dir)
def delete_obsolete_files(archive_name):
dir_name = 'dev_data/' + archive_name.replace('.tar.gz', '')
files = glob.glob(os.path.join(dir_name, '*.HB.DAT'))
for file in files:
os.remove(file)
files = glob.glob(os.path.join(dir_name, '*ADD.DAT'))
for file in files:
os.remove(file)
def normalised_metallicity_string(metallicity):
if not metallicity.startswith("0."):
metallicity = "0." + metallicity
metallicity_float = float(metallicity)
return "{:.4f}".format(metallicity_float)
def create_map_from_metallicity_to_archive_name(archive_names):
map = {}
metallicities = []
for archive_name in archive_names:
if archive_name.endswith('.tar.gz'):
match = re.search('Z(.*?)Y', archive_name)
metallicity = normalised_metallicity_string(match.group(1))
map[metallicity] = archive_name
metallicities.append(metallicity)
metallicities.sort()
return map, metallicities
def collect_mass_filenames(archive_name):
dir_name = 'dev_data/' + archive_name.replace('.tar.gz', '')
files = glob.glob(os.path.join(dir_name, '*.DAT'))
return files
def normalised_mass_string(mass):
mass_float = float(mass)
return "{:07.3f}".format(mass_float)
def create_map_from_mass_to_filename(filenames):
map = {}
masses = []
for file in filenames:
match = re.search('M(.*?).DAT', file)
mass = normalised_mass_string(match.group(1))
map[mass] = file
masses.append(mass)
masses.sort()
return map, masses
def clean_target_dir():
if not os.path.exists(TARGET_DIR):
os.makedirs(TARGET_DIR)
for file in os.listdir(TARGET_DIR):
os.remove(TARGET_DIR + file)
def generate_mod_file():
with open(TARGET_DIR + "mod.rs", 'w') as f:
f.write(MOD_TEMPLATE.format(URL=URL))
def generate_data_file(metallicities):
static_data = ""
access_array = ""
for index, metallicity in enumerate(metallicities):
variant_name = metallicity_variant_name(metallicity)
static_data += f"static {variant_name}_DATA: ParsecData = ParsecData::new({index});\n"
access_array += f"&{variant_name}_DATA,\n"
with open(TARGET_DIR + "data.rs", 'w') as f:
f.write(DATA_TEMPLATE.format(static_data=static_data,
array_size=len(metallicities),
access_array=access_array))
def mass_fraction_to_dex(mass_fraction):
Z_sun = 0.0122
return math.log10(mass_fraction / Z_sun)
def metallicity_variant_name(metallicity):
return "Z" + metallicity.replace(".", "_")
def generate_metallicity_file(metallicities, metallicity_to_archive_name):
mass_fraction_array = ""
dex_array = ""
names_array = ""
array_str = ""
archives_array = ""
to_mass_fraction_str = ""
to_dex_str = ""
for metallicity in metallicities:
variant_name = metallicity_variant_name(metallicity)
mass_fraction = float(metallicity)
dex = mass_fraction_to_dex(mass_fraction)
mass_fraction_array += f"{mass_fraction},\n"
dex_array += f"{dex},\n"
archive_name = metallicity_to_archive_name[metallicity]
names_array += f"\"{variant_name}\",\n"
archives_array += f"\"{archive_name}\",\n"
to_mass_fraction_str += f"Metallicity::{variant_name} => {mass_fraction},\n"
to_dex_str += f"Metallicity::{variant_name} => {dex},\n"
with open(TARGET_DIR + "metallicity.rs", 'w') as f:
f.write(METALLICITY_TEMPLATE.format(number_of_metallicities=len(metallicities),
mass_fraction_array=mass_fraction_array,
dex_array=dex_array,
names_array=names_array,
archives_array=archives_array,
metallicity_to_mass_fraction=to_mass_fraction_str,
metallicity_to_dex=to_dex_str))
def generate_masses_constant(metallicity, masses):
masses_str = f"const {metallicity_variant_name(metallicity)}_SORTED_MASSES: [f64; {len(masses)}] = ["
for mass in masses:
masses_str += f"{mass}, "
masses_str += "];"
return masses_str
def generate_filename_constant(metallicity, masses, mass_to_filename):
filenames_str = f"const {metallicity_variant_name(metallicity)}_FILENAMES: [&str; {len(masses)}] = ["
for mass in masses:
filename = os.path.basename(mass_to_filename[mass])
filenames_str += f"\"{filename}\", "
filenames_str += "];"
return filenames_str
def generate_masses_file(metallicities, metallicity_to_masses, metallicity_and_mass_to_filename):
sorted_masses = ""
filenames = ""
metallicity_to_masses_str = ""
metallicity_to_filenames_str = ""
for metallicity in metallicities:
masses = metallicity_to_masses[metallicity]
sorted_masses += generate_masses_constant(metallicity, masses) + "\n"
filenames += generate_filename_constant(metallicity, masses, metallicity_and_mass_to_filename[metallicity]) + "\n"
metallicity_name = metallicity_variant_name(metallicity)
metallicity_to_masses_str += f"&{metallicity_name}_SORTED_MASSES,\n"
metallicity_to_filenames_str += f"&{metallicity_name}_FILENAMES,\n"
with open(TARGET_DIR + "masses.rs", 'w') as f:
f.write(MASSES_TEMPLATE.format(sorted_masses=sorted_masses,
filenames=filenames,
number_of_metallicities=len(metallicities),
metallicity_to_masses=metallicity_to_masses_str,
metallicity_to_filenames=metallicity_to_filenames_str))
def main():
assure_dev_data_folder()
archive_names = collect_archive_names()
for archive_name in archive_names:
assure_archive_downloaded(archive_name)
assure_extracted(archive_name)
delete_obsolete_files(archive_name)
metallicity_to_archive_name, metallicities = create_map_from_metallicity_to_archive_name(archive_names)
metallicity_and_mass_to_filename = {}
metallicity_to_masses = {}
for metallicity, archive_name in metallicity_to_archive_name.items():
mass_filenames = collect_mass_filenames(archive_name)
mass_to_filename, masses = create_map_from_mass_to_filename(mass_filenames)
metallicity_and_mass_to_filename[metallicity] = mass_to_filename
metallicity_to_masses[metallicity] = masses
clean_target_dir()
generate_mod_file()
generate_data_file(metallicities)
generate_metallicity_file(metallicities, metallicity_to_archive_name)
generate_masses_file(metallicities, metallicity_to_masses, metallicity_and_mass_to_filename)
main()