Added simple unit tests for common methods. (054bcbe6) · Commits · Andrea Giannetti / swiss_army_knife_stable

etl/assets/commons/init.py

+5 −6

Original line number	Diff line number	Diff line
		@@ -3,6 +3,7 @@ import yaml
		import numpy as np
		import os
		import urllib.request
		from typing import List
		from astropy import units as u
		from assets.constants import (radmc_grid_map,
		radmc_coord_system_map,
		@@ -44,7 +45,7 @@ def compute_cartesian_coordinate_grid(indices: np.array,
		return indices * _physical_px_size_cm


		def extract_grid_metadata(config):
		def extract_grid_metadata(config: dict) -> dict:
		grid_config = config['grid']

		grid_metadata = grid_config.copy()
		@@ -52,8 +53,6 @@ def extract_grid_metadata(config):
		grid_metadata['coordinate_system'] = radmc_coord_system_map[grid_config['coordinate_system']]
		grid_metadata['grid_shape'] = get_grid_properties(grid_config=grid_config,
		keyword='steps')
		empty_grid = np.ones(grid_metadata['grid_shape'])

		grid_metadata['refpix'] = get_grid_properties(grid_config=grid_config,
		keyword='refpix')
		grid_metadata['grid_size'] = get_grid_properties(grid_config=grid_config,
		@@ -63,7 +62,7 @@ def extract_grid_metadata(config):

		grid_metadata['physical_px_size'] = get_physical_px_size(grid_metadata)

		indices = np.indices(empty_grid.shape).T - grid_metadata['refpix']
		indices = np.indices(grid_metadata['grid_shape']).T - grid_metadata['refpix']

		grid_metadata['coordinate_grid'] = compute_cartesian_coordinate_grid(indices=indices,
		physical_px_size=grid_metadata[
		@@ -89,11 +88,11 @@ def get_grid_edges(grid_metadata):
		def get_distance_matrix(grid_metadata, indices):
		distance_matrix = np.zeros(grid_metadata['grid_shape'])
		for axis_idx in range(len(grid_metadata['grid_shape'])):
		distance_matrix += (indices[:, :, :, axis_idx] * grid_metadata['physical_px_size'][axis_idx].value) ** 2
		distance_matrix += (indices[axis_idx, ...] * grid_metadata['physical_px_size'][axis_idx].to(u.cm).value) ** 2
		return np.sqrt(distance_matrix)


		def get_physical_px_size(grid_metadata):
		def get_physical_px_size(grid_metadata: dict) -> List[u.Quantity]:
		physical_px_size = []
		for axis_idx in range(len(grid_metadata['grid_shape'])):
		physical_px_size.append(

etl/tests/test_commons.py

+111 −2

Original line number	Diff line number	Diff line
		import numpy as np
		from assets.commons import (compute_power_law_radial_profile)
		import os
		from astropy import units as u
		from assets.commons import (compute_power_law_radial_profile,
		get_grid_properties,
		load_config_file,
		get_grid_edges,
		extract_grid_metadata,
		get_physical_px_size,
		compute_cartesian_coordinate_grid, get_distance_matrix)
		from unittest import TestCase


		def create_test_config(config_dict: dict,
		config_filename: str):
		with open(config_filename, 'w') as config_file:
		config_file.write('grid:\n')
		for key in config_dict:
		config_file.write(f' {key}: {config_dict[key]}\n')


		class Test(TestCase):
		def test_compute_power_law_radial_profile_flat(self):
		indices = np.indices([3, 3]) - 1
		@@ -25,3 +41,96 @@ class Test(TestCase):
		)
		self.assertTrue(np.array_equal(radial_profile, 1.0 / distance_matrix))

		def test_get_grid_properties_uniform(self):
		config_filename = os.path.join('test_files', 'config.yml')
		config_dict = {
		'dim1': {"size": 0.1, "unit": "pc", "steps": 5, "refpix": 2},
		}
		keywords = ['size', 'unit', 'steps', 'refpix']
		expected_result = {
		'size': [0.1] * 3,
		'unit': ['pc'] * 3,
		'steps': [5] * 3,
		'refpix': [2] * 3,
		}
		create_test_config(config_dict=config_dict,
		config_filename=config_filename)
		config = load_config_file(config_filename)
		for key in keywords:
		grid_properties = get_grid_properties(grid_config=config['grid'],
		keyword=key)
		self.assertListEqual(grid_properties, expected_result[key])

		def test_get_grid_properties(self):
		config_filename = os.path.join('test_files', 'config.yml')
		config_dict = {
		'dim1': {"size": 0.1, "unit": "pc", "steps": 5, "refpix": 2},
		'dim2': {"size": 0.2, "unit": "pc", "steps": 10, "refpix": 4.5},
		'dim3': {"size": 0.3, "unit": "pc", "steps": 15, "refpix": 7},
		}
		keywords = ['size', 'unit', 'steps', 'refpix']
		expected_result = {
		'size': [0.1, 0.2, 0.3],
		'unit': ['pc'] * 3,
		'steps': [5, 10, 15],
		'refpix': [2, 4.5, 7],
		}
		create_test_config(config_dict=config_dict,
		config_filename=config_filename)

		config = load_config_file(config_filename)
		for key in keywords:
		grid_properties = get_grid_properties(grid_config=config['grid'],
		keyword=key)
		self.assertListEqual(grid_properties, expected_result[key])

		def test_get_grid_edges(self):
		grid_metadata = {
		'grid_shape': [2, 2],
		'physical_px_size': [1 * u.cm, 1 * u.cm],
		'refpix': [0.5, 0.5]
		}
		expected_results_per_axis = np.array([-1, 0, 1])
		grid_edges = get_grid_edges(grid_metadata=grid_metadata)
		for axis_idx in range(len(grid_metadata['grid_shape'])):
		self.assertTrue(np.array_equal(expected_results_per_axis, grid_edges[:, axis_idx]))

		def test_get_physical_px_size(self):
		config_filename = os.path.join('test_files', 'config.yml')
		grid_size = 0.2
		npix = 5
		config_dict = {
		'grid_type': 'regular',
		'coordinate_system': 'cartesian',
		'central_density': '1e6',
		'density_unit': 'cm^-3',
		'dim1': {"size": grid_size, "unit": "pc", "steps": npix, "refpix": 2},
		'source_distance': '1000',
		'source_distance_unit': 'pc'
		}
		expected_result = [(grid_size * u.pc).to(u.cm) / npix]*3
		create_test_config(config_dict=config_dict,
		config_filename=config_filename)
		config = load_config_file(config_filename)
		grid_metadata = extract_grid_metadata(config=config)
		physical_px_size = get_physical_px_size(grid_metadata)
		self.assertListEqual(expected_result, physical_px_size)

		def test_compute_cartesian_coordinate_grid(self):
		indices = np.array([0, 1, 2])
		physical_px_size = [0.1 * u.pc]
		expected_results = indices * (physical_px_size[0]).to(u.cm).value
		computed_grid = compute_cartesian_coordinate_grid(indices=indices,
		physical_px_size=physical_px_size)
		self.assertTrue(np.array_equal(expected_results, computed_grid))

		def test_get_distance_matrix(self):
		grid_metadata = {
		'grid_shape': [2, 2],
		'physical_px_size': [1 * u.cm, 1 * u.cm]
		}
		indices = np.indices(grid_metadata['grid_shape'])
		expected_results = np.array([[0, 1], [1, np.sqrt(2)]])
		distance_matrix = get_distance_matrix(grid_metadata=grid_metadata,
		indices=indices)
		self.assertTrue(np.array_equal(expected_results, distance_matrix))