Skip to content

Commit

Permalink
slurmify couplings
Browse files Browse the repository at this point in the history
  • Loading branch information
xzackli committed Dec 12, 2023
1 parent 5494ebc commit 029697a
Show file tree
Hide file tree
Showing 7 changed files with 192 additions and 59 deletions.
20 changes: 20 additions & 0 deletions project/ana_cov_comp/paramfiles/1perlmutterdebug.slurm
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#! /bin/bash
#SBATCH -A mp107
#SBATCH -C cpu
#SBATCH -n 1 # node count
#SBATCH -c 128
#SBATCH -q debug
#SBATCH -t 0:30:00
#SBATCH --output=slurmoutput/R-%j.out

cd /pscratch/sd/x/xzackli/so/PSpipe_ana_cov_comp/project/ana_cov_comp
base_dir=/pscratch/sd/x/xzackli/so/envs/ps
pyenv_dir=${base_dir}/pyenv/perlmutter
export SOFIND_SYSTEM=perlmutter
export PYTHONPATH=$PYTHONPATH:${software_dir}
source ${pyenv_dir}/bin/activate

export OMP_NUM_THREADS=128

$1

19 changes: 19 additions & 0 deletions project/ana_cov_comp/paramfiles/1perlmutternode.slurm
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
#! /bin/bash
#SBATCH -A mp107
#SBATCH -C cpu
#SBATCH -n 1 # node count
#SBATCH -c 128
#SBATCH -q regular
#SBATCH -t 1:00:00
#SBATCH --output=slurmoutput/R-%j.out

cd /pscratch/sd/x/xzackli/so/PSpipe_ana_cov_comp/project/ana_cov_comp
base_dir=/pscratch/sd/x/xzackli/so/envs/ps
pyenv_dir=${base_dir}/pyenv/perlmutter
export SOFIND_SYSTEM=perlmutter
export PYTHONPATH=$PYTHONPATH:${software_dir}
source ${pyenv_dir}/bin/activate

export OMP_NUM_THREADS=128

$1
8 changes: 4 additions & 4 deletions project/ana_cov_comp/paramfiles/cov_dr6_v4_20231128.dict
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
data_dir = '/scratch/gpfs/zatkins/data/simonsobs/PSpipe/project/ana_cov_comp/cov_dr6_v4_20231128/'
data_dir = '/pscratch/sd/x/xzackli/so/data/dr6v4/'

surveys = ["dr6"]
arrays_dr6 = {'pa4': ['f220'], 'pa5': ['f090', 'f150'], 'pa6': ['f090', 'f150']}
Expand All @@ -11,7 +11,7 @@ multistep_path = data_dir
mcms_dir = data_dir + 'mcms'

binned_mcm = False
lmax = 5400
lmax = 8500
use_toeplitz_mcm = False
binning_file = data_dir + "binning/BIN_ACTPOL_50_4_SC_large_bin_at_low_ell"
niter = 0
Expand Down Expand Up @@ -78,7 +78,7 @@ edge_skip_rescale = 1
cross_link_threshold = 0.97
n_med_ivar = 3

window_dir = '/scratch/gpfs/zatkins/data/simonsobs/mnms/masks/tlouis/ACT/'
window_dir = '/pscratch/sd/x/xzackli/so/data/windows/'

window_T_dr6_pa4_f150 = window_dir + "window_dr6_pa4_f150_baseline.fits"
window_pol_dr6_pa4_f150 = window_dir + "window_dr6_pa4_f150_baseline.fits"
Expand Down Expand Up @@ -116,7 +116,7 @@ window_T_dr6_pa6_f150_alias = "w_dr6_pa6_f150"
window_pol_dr6_pa6_f150_alias = "w_dr6_pa6_f150"

# maps
map_dir = '/scratch/gpfs/ACT/dr6v4/maps/dr6v4_20230316/release/'
map_dir = '/pscratch/sd/x/xzackli/so/data/dr6v4/maps/'

n_splits_dr6 = 4
src_free_maps_dr6 = True
Expand Down
50 changes: 24 additions & 26 deletions project/ana_cov_comp/python/get_2pt_coupling_matrices.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,25 +98,24 @@
# the coupling
if (ewin_name1, ewin_name2, spintype) not in canonized_combos:
canonized_combos[(ewin_name1, ewin_name2, spintype)] = [(field_info1, field_info2, spintype)]
assert not os.path.isfile(coupling_fn), \
f'{coupling_fn} exists but we should not yet have produced it in loop'
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name2}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)
if not os.path.isfile(coupling_fn):
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name2}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)

# if we have already gotten to this canonized window pair, add the
# fields to the tracked list and ensure we've already calculated the
# coupling
else:
canonized_combos[(ewin_name1, ewin_name2, spintype)].append((field_info1, field_info2, spintype))
assert os.path.isfile(coupling_fn), \
f'{coupling_fn} does not exist but we should have produced it in loop'
# assert os.path.isfile(coupling_fn), \
# f'{coupling_fn} does not exist but we should have produced it in loop'
continue

# N N
Expand Down Expand Up @@ -144,25 +143,24 @@
# the coupling
if (ewin_name1, ewin_name2, spintype) not in canonized_combos:
canonized_combos[(ewin_name1, ewin_name2, spintype)] = [(field_info1, field_info2, spintype)]
assert not os.path.isfile(coupling_fn), \
f'{coupling_fn} exists but we should not yet have produced it in loop'
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name2}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)
if not os.path.isfile(coupling_fn):
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name2}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)

# if we have already gotten to this canonized window pair, add the
# fields to the tracked list and ensure we've already calculated the
# coupling
else:
canonized_combos[(ewin_name1, ewin_name2, spintype)].append((field_info1, field_info2, spintype))
assert os.path.isfile(coupling_fn), \
f'{coupling_fn} does not exist but we should have produced it in loop'
# assert os.path.isfile(coupling_fn), \
# f'{coupling_fn} does not exist but we should have produced it in loop'
continue

np.save(f'{couplings_dir}/canonized_couplings_2pt_combos.npy', canonized_combos)
2 changes: 1 addition & 1 deletion project/ana_cov_comp/python/get_2pt_ewin_alms.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
if field_info not in field_infos:
field_infos.append(field_info)
else:
raise ValueError(f'{field_info=} is not unique')
raise ValueError(f'{field_info} is not unique')

ewin_info_s = psc.get_ewin_info_from_field_info(field_info, d, mode='w')
if ewin_info_s not in ewin_infos:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
# we define the canon by the windows order. we first build the fields,
# then use a mapping from fields to windows to build the canonical
# windows
recipe = {'S_only':[], 'SxN':[], 'N_only':[]}
field_infos = []
ewin_infos = []
for sv1 in surveys:
Expand Down Expand Up @@ -123,21 +124,29 @@
f'{coupling_fn} exists but we should not yet have produced it in loop'
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)
recipe['S_only'].append({
'spintype': spintype,
'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy',
'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy',
'lmax':lmax,
'input_alm': True,
'coupling_fn': coupling_fn
})
# w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
# w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
# coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
# lmax=lmax, input_alm=True,
# l_exact=l_exact, l_toep=l_toep,
# l_band=l_band)
# np.save(coupling_fn, coupling)

# if we have already gotten to this canonized window pair, add the
# fields to the tracked list and ensure we've already calculated the
# coupling
else:
canonized_combos[(ewin_name1, ewin_name2, ewin_name3, ewin_name4, spintype)].append((field_info1, field_info2, field_info3, field_info4, spintype))
assert os.path.isfile(coupling_fn), \
f'{coupling_fn} does not exist but we should have produced it in loop'
# assert os.path.isfile(coupling_fn), \
# f'{coupling_fn} does not exist but we should have produced it in loop'
continue

# S S N N and N N S S
Expand Down Expand Up @@ -173,21 +182,29 @@
f'{coupling_fn} exists but we should not yet have produced it in loop'
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)
recipe['SxN'].append({
'spintype': spintype,
'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy',
'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy',
'lmax':lmax,
'input_alm': True,
'coupling_fn': coupling_fn
})
# w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
# w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
# coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
# lmax=lmax, input_alm=True,
# l_exact=l_exact, l_toep=l_toep,
# l_band=l_band)
# np.save(coupling_fn, coupling)

# if we have already gotten to this canonized window pair, add the
# fields to the tracked list and ensure we've already calculated the
# coupling
else:
canonized_combos[(ewin_name1, ewin_name2, ewin_name3, ewin_name4, spintype)].append((field_info1, field_info2, field_info3, field_info4, spintype))
assert os.path.isfile(coupling_fn), \
f'{coupling_fn} does not exist but we should have produced it in loop'
# assert os.path.isfile(coupling_fn), \
# f'{coupling_fn} does not exist but we should have produced it in loop'
continue

# N N N N
Expand Down Expand Up @@ -223,21 +240,36 @@
f'{coupling_fn} exists but we should not yet have produced it in loop'
log.info(os.path.splitext(os.path.basename(coupling_fn))[0])

w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)

recipe['N_only'].append({
'spintype': spintype,
'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy',
'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy',
'lmax':lmax,
'input_alm': True,
'coupling_fn': coupling_fn
})

# w1 = np.load(f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy')
# w2 = np.load(f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy')
# coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
# lmax=lmax, input_alm=True,
# l_exact=l_exact, l_toep=l_toep,
# l_band=l_band)
# np.save(coupling_fn, coupling)

# if we have already gotten to this canonized window pair, add the
# fields to the tracked list and ensure we've already calculated the
# coupling
else:
canonized_combos[(ewin_name1, ewin_name2, ewin_name3, ewin_name4, spintype)].append((field_info1, field_info2, field_info3, field_info4, spintype))
assert os.path.isfile(coupling_fn), \
f'{coupling_fn} does not exist but we should have produced it in loop'
# assert os.path.isfile(coupling_fn), \
# f'{coupling_fn} does not exist but we should have produced it in loop'
continue

np.save(f'{couplings_dir}/canonized_couplings_4pt_combos.npy', canonized_combos)
total_coupling_number = len(recipe['S_only']) + len(recipe['SxN']) + len(recipe['N_only'])
recipe['total'] = total_coupling_number
np.save(f'{couplings_dir}/4pt_recipe.npy', recipe)
log.info(f"Total 4pt Couplings: {total_coupling_number}")
np.save(f'{couplings_dir}/4pt_recipe.npy', recipe)
np.save(f'{couplings_dir}/canonized_couplings_4pt_combos.npy', canonized_combos)
64 changes: 64 additions & 0 deletions project/ana_cov_comp/python/get_4pt_coupling_result.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
"""
This script reads the recipe from get_4pt_coupling_matrices_recipe.py and then cooks!
# For production on all, do something like
alias shortjob="sbatch paramfiles/1perlmutternode.slurm $1" # short QOS
for i in {0..57}; do \
shortjob "srun --ntasks 1 --cpus-per-task 128 --cpu-bind=cores python -u \
python/get_4pt_coupling_result.py paramfiles/cov_dr6_v4_20231128.dict $((50*i)) $((50*i+50))"
done
# For testing only:
alias testjob="sbatch paramfiles/1perlmutterdebug.slurm $1" # test QOS
testjob "srun --ntasks 1 --cpus-per-task 128 --cpu-bind=cores python -u \
python/get_4pt_coupling_result.py paramfiles/cov_dr6_v4_20231128.dict 1 2"
"""

import sys
import numpy as np
from pspipe_utils import log, covariance as psc
from pspy import so_dict, so_map, so_mcm, pspy_utils
from itertools import product, combinations_with_replacement as cwr
import os

d = so_dict.so_dict()
d.read_from_file(sys.argv[1])

log = log.get_logger(**d)

ewin_alms_dir = d['ewin_alms_dir']
couplings_dir = d['couplings_dir']
pspy_utils.create_directory(couplings_dir)

surveys = d['surveys']
arrays = {sv: d[f'arrays_{sv}'] for sv in surveys}

if d['use_toeplitz_cov'] == True:
log.info('we will use the toeplitz approximation')
l_exact, l_band, l_toep = 800, 2000, 2750
else:
l_exact, l_band, l_toep = None, None, None

lmax = d['lmax']


recipe = np.load(f'{couplings_dir}/4pt_recipe.npy', allow_pickle=True)[()]

unrolled = recipe['S_only'] + recipe['SxN'] + recipe['N_only']
total_couplings = len(unrolled)

if len(sys.argv) == 4:
log.info(f"computing only the covariance matrices : " +
f"{int(sys.argv[2])}:{int(sys.argv[3])} of {total_couplings}")
unrolled = unrolled[int(sys.argv[2]):int(sys.argv[3])]

for k in unrolled:
spintype, filename1, filename2, coupling_fn = k['spintype'], k['w1'], k['w2'], k['coupling_fn']
w1 = np.load(filename1)
w2 = np.load(filename2)
log.info(f"{spintype} {coupling_fn} from {filename1} {filename2}")
coupling = so_mcm.coupling_block(spintype, win1=w1, win2=w2,
lmax=lmax, input_alm=True,
l_exact=l_exact, l_toep=l_toep,
l_band=l_band)
np.save(coupling_fn, coupling)

0 comments on commit 029697a

Please sign in to comment.