spacetelescope / PASTIS

Algorithm for analytical contrast predictions of coronagraphs on segmented telescopes
BSD 3-Clause "New" or "Revised" License
8 stars 4 forks source link

Minor improvement in SM generation #159

Closed ivalaginja closed 1 year ago

ivalaginja commented 1 year ago

I am adding some simplifications and clarifying comments. End goal is to speed up the creation of the segmented mirrors in a simulator, see #160.

The script below speeds up from a runtime of 2.43608 sec to 2.31140 sec in the test script below, which is about 5% faster. This is not much but shaves off a couple of minutes on the really large mirrors.

Testing script used for this work:

import os
import time
import hcipy

from pastis.config import CONFIG_PASTIS
from pastis.simulators.scda_telescopes import HexRingAPLC
from pastis import util

NUM_RINGS = 2
optics_input = os.path.join(util.find_repo_location(), 'data', 'SCDA')
sampling = CONFIG_PASTIS.getfloat('HexRingTelescope', 'sampling')
sim = HexRingAPLC(optics_input, NUM_RINGS, sampling)

t1 = time.perf_counter()

n_zernikes = 2
seg_evaluated = sim._create_evaluated_segment_grid()

# Create a single segment influence function with all Zernikes n_zernikes
# This allows us to have an initial mode basis to extend with all the other segments later on
first_seg = 0  # Create this first influence function on the first segment only
local_zernike_basis = hcipy.mode_basis.make_zernike_basis(n_zernikes,
                                                          sim.segment_circumscribed_diameter,
                                                          sim.pupil_grid.shifted(-sim.seg_pos[first_seg]),
                                                          starting_mode=1, radial_cutoff=False)

# For all Zernikes on this first segment, cut them to the actual segment support
for zernike_num in range(0, n_zernikes):
    local_zernike_basis._transformation_matrix[:, zernike_num] *= seg_evaluated[first_seg]

# Expand the basis of influence functions from one segment to all segments
for seg_num in range(1, sim.nseg):
    local_zernike_basis_tmp = hcipy.mode_basis.make_zernike_basis(n_zernikes,
                                                                  sim.segment_circumscribed_diameter,
                                                                  sim.pupil_grid.shifted(-sim.seg_pos[seg_num]),
                                                                  starting_mode=1, radial_cutoff=False)
    # Cut to segment support
    for zernike_num in range(0, n_zernikes):
        local_zernike_basis_tmp._transformation_matrix[:, zernike_num] *= seg_evaluated[seg_num]
    local_zernike_basis.extend(local_zernike_basis_tmp)  # extend our basis with this new segment

t2 = time.perf_counter()

print(local_zernike_basis.transformation_matrix.shape)
# hcipy.write_mode_basis(local_zernike_basis, 'basis.asdf')
print(f'Elapsed time for SM generation: {t2 - t1}')