Skip to content

Commit

Permalink
Resoring functions requried for celeri_solve.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@brendanjmeade
brendanjmeade committed Jul 10, 2024
1 parent b278a38 commit 1d08184
Showing 1 changed file with 288 additions and 0 deletions.
288 changes: 288 additions & 0 deletions celeri/celeri.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1455,6 +1455,47 @@ def get_index_no_meshes(assembly, station, block):
return index


def get_processed_data_structures(command):
# NOTE: Used in celeri_solve.py
data = addict.Dict()
assembly = addict.Dict()
operators = addict.Dict()
(
data.segment,
data.block,
data.meshes,
data.station,
data.mogi,
data.sar,
) = read_data(command)
data.station = process_station(data.station, command)
data.segment = process_segment(data.segment, command, data.meshes)
data.sar = process_sar(data.sar, command)
data.closure, data.block = assign_block_labels(
data.segment, data.station, data.block, data.mogi, data.sar
)
operators.meshes = [addict.Dict()] * len(data.meshes)
assembly = merge_geodetic_data(
assembly, data.station, data.sar
) # TODO: Not sure this works correctly

# Quick input plot
if bool(command.plot_input_summary):
plot_input_summary(
command,
data.segment,
data.station,
data.block,
data.meshes,
data.mogi,
data.sar,
lon_range=command.lon_range,
lat_range=command.lat_range,
quiver_scale=command.quiver_scale,
)
return data, assembly, operators


##############################################################################################################################
# #
# #
Expand Down Expand Up @@ -4386,6 +4427,253 @@ def post_process_estimation_hmatrix(
estimation_hmatrix.north_vel_tde = estimation_hmatrix.vel_tde[1::2]


def build_and_solve_dense(command, assembly, operators, data):
# NOTE: Used in celeri_solve.py
logger.info("build_and_solve_dense")

# Get all elastic operators for segments and TDEs
get_elastic_operators(operators, data.meshes, data.segment, data.station, command)

# Get TDE smoothing operators
get_all_mesh_smoothing_matrices(data.meshes, operators)

# Get non-elastic operators
operators.rotation_to_velocities = get_rotation_to_velocities_partials(
data.station, data.block.shape[0]
)
operators.global_float_block_rotation = get_global_float_block_rotation_partials(
data.station
)
assembly, operators.block_motion_constraints = get_block_motion_constraints(
assembly, data.block, command
)
assembly, operators.slip_rate_constraints = get_slip_rate_constraints(
assembly, data.segment, data.block, command
)
operators.rotation_to_slip_rate = get_rotation_to_slip_rate_partials(
data.segment, data.block
)
(
operators.block_strain_rate_to_velocities,
strain_rate_block_index,
) = get_block_strain_rate_to_velocities_partials(
data.block, data.station, data.segment
)
operators.mogi_to_velocities = get_mogi_to_velocities_partials(
data.mogi, data.station, command
)
get_tde_slip_rate_constraints(data.meshes, operators)

# Direct solve dense linear system
logger.info("Start: Dense assemble and solve")
start_solve_time = timeit.default_timer()
index, estimation = assemble_and_solve_dense(
command, assembly, operators, data.station, data.block, data.meshes
)
end_solve_time = timeit.default_timer()
logger.success(
f"Finish: Dense assemble and solve: {end_solve_time - start_solve_time:0.2f} seconds for solve"
)

post_process_estimation(estimation, operators, data.station, index)

write_output(
command, estimation, data.station, data.segment, data.block, data.meshes
)

if bool(command.plot_estimation_summary):
plot_estimation_summary(
command,
data.segment,
data.station,
data.meshes,
estimation,
lon_range=command.lon_range,
lat_range=command.lat_range,
quiver_scale=command.quiver_scale,
)

return estimation, operators, index


def build_and_solve_dense_no_meshes(command, assembly, operators, data):
# NOTE: Used in celeri_solve.py
logger.info("build_and_solve_dense_no_meshes")

# Get all elastic operators for segments and TDEs
get_elastic_operators(operators, data.meshes, data.segment, data.station, command)

operators.rotation_to_velocities = get_rotation_to_velocities_partials(
data.station, data.block.shape[0]
)
operators.global_float_block_rotation = get_global_float_block_rotation_partials(
data.station
)
assembly, operators.block_motion_constraints = get_block_motion_constraints(
assembly, data.block, command
)
assembly, operators.slip_rate_constraints = get_slip_rate_constraints(
assembly, data.segment, data.block, command
)
operators.rotation_to_slip_rate = get_rotation_to_slip_rate_partials(
data.segment, data.block
)
(
operators.block_strain_rate_to_velocities,
strain_rate_block_index,
) = get_block_strain_rate_to_velocities_partials(
data.block, data.station, data.segment
)
operators.mogi_to_velocities = get_mogi_to_velocities_partials(
data.mogi, data.station, command
)

# Blocks only operator
index = get_index_no_meshes(assembly, data.station, data.block)

# TODO: Clean up!
logger.error(operators.keys())

operator_block_only = get_full_dense_operator_block_only(operators, index)
# weighting_vector = get_weighting_vector(command, data.station, data.meshes, index)
weighting_vector = get_weighting_vector_no_meshes(command, data.station, index)
data_vector = get_data_vector(assembly, index)
weighting_vector_block_only = weighting_vector[0 : operator_block_only.shape[0]]

# Solve the overdetermined linear system using only a weighting vector rather than matrix
estimation = addict.Dict()
estimation.operator = operator_block_only
estimation.weighting_vector = weighting_vector_block_only

estimation.state_covariance_matrix = np.linalg.inv(
operator_block_only.T * weighting_vector_block_only @ operator_block_only
)
estimation.state_vector = (
estimation.state_covariance_matrix
@ operator_block_only.T
* weighting_vector_block_only
@ data_vector[0 : weighting_vector_block_only.size]
)

# Post-processing

estimation.predictions = estimation.operator @ estimation.state_vector
estimation.vel = estimation.predictions[0 : 2 * index.n_stations]
estimation.east_vel = estimation.vel[0::2]
estimation.north_vel = estimation.vel[1::2]

# Estimate slip rate uncertainties
estimation.slip_rate_sigma = np.sqrt(
np.diag(
operators.rotation_to_slip_rate
@ estimation.state_covariance_matrix[
0 : 3 * index.n_blocks, 0 : 3 * index.n_blocks
]
@ operators.rotation_to_slip_rate.T
)
) # I don't think this is correct because for the case when there is a rotation vector a priori
estimation.strike_slip_rate_sigma = estimation.slip_rate_sigma[0::3]
estimation.dip_slip_rate_sigma = estimation.slip_rate_sigma[1::3]
estimation.tensile_slip_rate_sigma = estimation.slip_rate_sigma[2::3]

# Calculate mean squared residual velocity
estimation.east_vel_residual = estimation.east_vel - data.station.east_vel
estimation.north_vel_residual = estimation.north_vel - data.station.north_vel

# Extract segment slip rates from state vector
estimation.slip_rates = (
operators.rotation_to_slip_rate
@ estimation.state_vector[0 : 3 * index.n_blocks]
)
estimation.strike_slip_rates = estimation.slip_rates[0::3]
estimation.dip_slip_rates = estimation.slip_rates[1::3]
estimation.tensile_slip_rates = estimation.slip_rates[2::3]

# Calculate rotation only velocities
estimation.vel_rotation = (
operators.rotation_to_velocities[index.station_row_keep_index, :]
@ estimation.state_vector[0 : 3 * index.n_blocks]
)
estimation.east_vel_rotation = estimation.vel_rotation[0::2]
estimation.north_vel_rotation = estimation.vel_rotation[1::2]

# Calculate fully locked segment velocities
estimation.vel_elastic_segment = (
operators.rotation_to_slip_rate_to_okada_to_velocities[
index.station_row_keep_index, :
]
@ estimation.state_vector[0 : 3 * index.n_blocks]
)
estimation.east_vel_elastic_segment = estimation.vel_elastic_segment[0::2]
estimation.north_vel_elastic_segment = estimation.vel_elastic_segment[1::2]

# TODO: Calculate block strain rate velocities
estimation.east_vel_block_strain_rate = np.zeros(len(data.station))
estimation.north_vel_block_strain_rate = np.zeros(len(data.station))

# # Get all elastic operators for segments and TDEs
# get_elastic_operators(operators, data.meshes, data.segment, data.station, command)

# # Get TDE smoothing operators
# get_all_mesh_smoothing_matrices(data.meshes, operators)

# # Get non-elastic operators
# operators.rotation_to_velocities = get_rotation_to_velocities_partials(data.station, data.block.shape[0])
# operators.global_float_block_rotation = get_global_float_block_rotation_partials(
# data.station
# )
# assembly, operators.block_motion_constraints = get_block_motion_constraints(
# assembly, data.block, command
# )
# assembly, operators.slip_rate_constraints = get_slip_rate_constraints(
# assembly, data.segment, data.block, command
# )
# operators.rotation_to_slip_rate = get_rotation_to_slip_rate_partials(
# data.segment, data.block
# )
# (
# operators.block_strain_rate_to_velocities,
# strain_rate_block_index,
# ) = get_block_strain_rate_to_velocities_partials(
# data.block, data.station, data.segment
# )
# operators.mogi_to_velocities = get_mogi_to_velocities_partials(
# data.mogi, data.station, command
# )
# get_tde_slip_rate_constraints(data.meshes, operators)

# # Direct solve dense linear system
# logger.info("Start: Dense assemble and solve")
# start_solve_time = timeit.default_timer()
# index, estimation = assemble_and_solve_dense(
# command, assembly, operators, data.station, data.block, data.meshes
# )
# end_solve_time = timeit.default_timer()
# logger.success(
# f"Finish: Dense assemble and solve: {end_solve_time - start_solve_time:0.2f} seconds for solve"
# )

# post_process_estimation(estimation, operators, data.station, index)

write_output(
command, estimation, data.station, data.segment, data.block, data.meshes
)

if bool(command.plot_estimation_summary):
plot_estimation_summary(
command,
data.segment,
data.station,
data.meshes,
estimation,
lon_range=command.lon_range,
lat_range=command.lat_range,
quiver_scale=command.quiver_scale,
)

return estimation, operators, index


#########################################################################################
# #
# #
Expand Down

0 comments on commit 1d08184

Please sign in to comment.