Skip to content

Commit

Permalink
Improve tides analysis step
Browse files Browse the repository at this point in the history 
 - Run extract_HC in parallel

 - Add list of constituents to analyze to config file
  • Loading branch information
@sbrus89
sbrus89 committed Dec 3, 2024
1 parent db6fda2 commit bd62b1c
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 85 deletions.
190 changes: 105 additions & 85 deletions compass/ocean/tests/tides/analysis/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import os
import subprocess

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import netCDF4
import numpy as np
from matplotlib import colormaps
from mpas_tools.logging import check_call

from compass.step import Step

Expand Down Expand Up @@ -42,7 +42,8 @@ def __init__(self, test_case):

self.harmonic_analysis_file = 'harmonicAnalysis.nc'
self.grid_file = 'initial_state.nc'
self.constituents = ['k1', 'k2', 'm2', 'n2', 'o1', 'p1', 'q1', 's2']
self.constituents_valid = ['k1', 'k2', 'm2', 'n2',
'o1', 'p1', 'q1', 's2']

self.add_input_file(
filename=self.harmonic_analysis_file,
Expand All @@ -60,6 +61,8 @@ def setup(self):
config = self.config
self.tpxo_version = config.get('tides', 'tpxo_version')

self.constituents = config.getlist('tides', 'constituents')

os.makedirs(f'{self.work_dir}/TPXO_data', exist_ok=True)
if self.tpxo_version == 'TPXO9':
for constituent in self.constituents:
Expand Down Expand Up @@ -90,72 +93,79 @@ def setup(self):
target='TPXO8/grid_tpxo8_atlas_30_v1',
database='tides')

def write_coordinate_file(self, idx):
def write_coordinate_file(self, nchunks):
"""
Write mesh coordinates for TPXO extraction
"""

# Read in mesh
grid_nc = netCDF4.Dataset(self.grid_file, 'r')
lon_grid = np.degrees(grid_nc.variables['lonCell'][idx])
lat_grid = np.degrees(grid_nc.variables['latCell'][idx])
nCells = len(lon_grid)
lon_grid = np.degrees(grid_nc.variables['lonCell'][:])
lat_grid = np.degrees(grid_nc.variables['latCell'][:])

lon_grid_split = np.array_split(lon_grid, nchunks)
lat_grid_split = np.array_split(lat_grid, nchunks)

# Write coordinate file for OTPS2
f = open('lat_lon', 'w')
for i in range(nCells):
f.write(f'{lat_grid[i]} {lon_grid[i]} \n')
f.close()
for chunk in range(nchunks):
f = open(f'inputs/lat_lon_{chunk}', 'w')
for i in range(lon_grid_split[chunk].size):
f.write(f'{lat_grid_split[chunk][i]} '
f'{lon_grid_split[chunk][i]} \n')
f.close()

def setup_otps2(self):
def setup_otps2(self, nchunks):
"""
Write input files for TPXO extraction
"""

for con in self.constituents:
print(f'setup {con}')

# Lines for the setup_con files
lines = [{'inp': f'inputs/Model_atlas_{con}',
'comment': '! 1. tidal model control file'},
{'inp': 'lat_lon',
'comment': '! 2. latitude/longitude/<time> file'},
{'inp': 'z',
'comment': '! 3. z/U/V/u/v'},
{'inp': con,
'comment': '! 4. tidal constituents to include'},
{'inp': 'AP',
'comment': '! 5. AP/RI'},
{'inp': 'oce',
'comment': '! 6. oce/geo'},
{'inp': '1',
'comment': '! 7. 1/0 correct for minor constituents'},
{'inp': f'outputs/{con}.out',
'comment': '! 8. output file (ASCII)'}]

# Create directory for setup_con and Model_atlas_con files
if not os.path.exists('inputs'):
os.mkdir('inputs')

# Write the setup_con file
f = open(f'inputs/{con}_setup', 'w')
for line in lines:
spaces = 28 - len(line['inp'])
f.write(line['inp'] + spaces * ' ' + line['comment'] + '\n')
f.close()

# Write the Model_atlas_con file
f = open(f'inputs/Model_atlas_{con}', 'w')
f.write(f'TPXO_data/h_{con}_tpxo\n')
f.write(f'TPXO_data/u_{con}_tpxo\n')
f.write('TPXO_data/grid_tpxo')
f.close()

# Create directory for the con.out files
if not os.path.exists('outputs'):
os.mkdir('outputs')

def run_otps2(self):
for chunk in range(nchunks):

# Lines for the setup_con files
lines = [{'inp': f'inputs/Model_atlas_{con}',
'comment': '! 1. tidal model control file'},
{'inp': f'inputs/lat_lon_{chunk}',
'comment': '! 2. latitude/longitude/<time> file'},
{'inp': 'z',
'comment': '! 3. z/U/V/u/v'},
{'inp': con,
'comment': '! 4. tidal constituents to include'},
{'inp': 'AP',
'comment': '! 5. AP/RI'},
{'inp': 'oce',
'comment': '! 6. oce/geo'},
{'inp': '1',
'comment':
'! 7. 1/0 correct for minor constituents'},
{'inp': f'outputs/{con}_{chunk}.out',
'comment': '! 8. output file (ASCII)'}]

# Create directory for setup_con and Model_atlas_con files
if not os.path.exists('inputs'):
os.mkdir('inputs')

# Write the setup_con file
f = open(f'inputs/{con}_setup_{chunk}', 'w')
for line in lines:
spaces = (28 - len(line['inp'])) * ' '
f.write(f"{line['inp']}{spaces}{line['comment']}\n")
f.close()

# Write the Model_atlas_con file
f = open(f'inputs/Model_atlas_{con}', 'w')
f.write(f'TPXO_data/h_{con}_tpxo\n')
f.write(f'TPXO_data/u_{con}_tpxo\n')
f.write('TPXO_data/grid_tpxo')
f.close()

# Create directory for the con.out files
if not os.path.exists('outputs'):
os.mkdir('outputs')

def run_otps2(self, nchunks):
"""
Perform TPXO extraction
"""
Expand All @@ -164,35 +174,46 @@ def run_otps2(self):
for con in self.constituents:
print('')
print(f'run {con}')
check_call(f'extract_HC < inputs/{con}_setup',
logger=self.logger, shell=True)

def read_otps2_output(self, idx):
commands = []
for chunk in range(nchunks):
commands.append(f'extract_HC < inputs/{con}_setup_{chunk}')

processes = [subprocess.Popen(cmd, shell=True) for cmd in commands]

for p in processes:
p.wait()

def read_otps2_output(self, nchunks):
"""
Read TPXO extraction output
"""

start = idx[0]
for con in self.constituents:

f = open(f'outputs/{con}.out', 'r')
lines = f.read().splitlines()
for i, line in enumerate(lines[3:]):
line_sp = line.split()
if line_sp[2] != '*************':
val = float(line_sp[2])
self.mesh_AP[con]['amp'][start + i] = val
else:
self.mesh_AP[con]['amp'][start + i] = -9999

if line_sp[3] != 'Site':
val = float(line_sp[3])
if val < 0:
val = val + 360.0
self.mesh_AP[con]['phase'][start + i] = val

else:
self.mesh_AP[con]['phase'][start + i] = -9999
i = 0
for chunk in range(nchunks):

f = open(f'outputs/{con}_{chunk}.out', 'r')
lines = f.read().splitlines()
for line in lines[3:]:
line_sp = line.split()
if line_sp[2] != '*************':
val = float(line_sp[2])
self.mesh_AP[con]['amp'][i] = val
else:
self.mesh_AP[con]['amp'][i] = -9999

if line_sp[3] != 'Site':
val = float(line_sp[3])
if val < 0:
val = val + 360.0
self.mesh_AP[con]['phase'][i] = val

else:
self.mesh_AP[con]['phase'][i] = -9999

i = i + 1

def append_tpxo_data(self):
"""
Expand Down Expand Up @@ -283,7 +304,7 @@ def plot(self):
depth[:] = data_nc.variables['bottomDepth'][:]
area[:] = data_nc.variables['areaCell'][:]

constituent_list = ['K1', 'M2', 'N2', 'O1', 'S2']
constituent_list = [x.upper() for x in self.constituents]

# Use these to fix up the plots
subplot_ticks = [[np.linspace(0, 0.65, 10), np.linspace(0, 0.65, 10),
Expand Down Expand Up @@ -429,16 +450,16 @@ def run(self):
Run this step of the test case
"""

self.constituents = self.config.getlist('tides', 'constituents')

# Check if TPXO values aleady exist in harmonic_analysis.nc
self.check_tpxo_data()

# Setup input files for TPXO extraction
self.setup_otps2()

# Setup chunking for TPXO extraction with large meshes
indices = np.arange(self.nCells)
nchunks = np.ceil(self.nCells / 200000)
index_chunks = np.array_split(indices, nchunks)
nchunks = 128

# Setup input files for TPXO extraction
self.setup_otps2(nchunks)

# Initialize data structure for TPXO values
self.mesh_AP = {}
Expand All @@ -447,10 +468,9 @@ def run(self):
'phase': np.zeros((self.nCells))}

# Extract TPXO values
for idx in index_chunks:
self.write_coordinate_file(idx)
self.run_otps2()
self.read_otps2_output(idx)
self.write_coordinate_file(nchunks)
self.run_otps2(nchunks)
self.read_otps2_output(nchunks)

# Inject TPXO values
self.append_tpxo_data()
Expand Down
3 changes: 3 additions & 0 deletions compass/ocean/tests/tides/mesh/icos7/icos7.cfg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ forward_threads = 1
# TPXO version for validation
tpxo_version = TPXO9

# constituents to validate
constituents = k1, k2, m2, n2, o1, p1, q1, s2


# config options related to remapping topography to an MPAS-Ocean mesh
[remap_topography]
Expand Down
3 changes: 3 additions & 0 deletions compass/ocean/tests/tides/mesh/vr45to5/vr45to5.cfg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,9 @@ forward_threads = 1
# TPXO version for validation
tpxo_version = TPXO9

# constituents to validate
constituents = k1, k2, m2, n2, o1, p1, q1, s2

# config options related to remapping topography to an MPAS-Ocean mesh
[remap_topography]

Expand Down

0 comments on commit bd62b1c

Please sign in to comment.