Steps needed to invert model output

[cspencerjones]

This is my first attempt at writing down the steps required to get to a place where we can invert model output to get the transport matrix.
Comments welcome!

• Replace cost function: tracer plus tracer control vector should equal model output of tracer (22 in the notes)
• Check that with noisy input we get back the model output of tracer
• Tracer control vector should be minimized (first term of 18)
• The vector f (and not the lu factorized version of A) should be an input to the cost function
• It then probably makes sense to add tracer conservation (17 in the notes) to the cost function
• Find the Lagrange multiplier for the tracer (26)
• Test to see if tracer is conserved more accurately after optimization (Edit: there is actually no reason for this to be true at this point)
• A second control vector u_f should also be an input to the cost function
• Changes in u_f should be minimized (second term of 18)
• Add constraint that f cannot become negative (21)
• Add Lagrange multiplier from first term of (27)
• Add mass conservation to cost function (19) and to Lagrange multiplier
• Test to see if mass is conserved more accurately after optimization
• Add geostrophy to cost function (20) and to Lagrange multiplier
• Test to see if geostrophic balance is conserved more accurately

[cspencerjones]

I have been playing around here:

TMI.jl/src/TMI.jl

Lines 2219 to 2295 in d097182

	function costfunction_gridded_model(convec::Vector{T},non_zero_indices,b::BoundaryCondition{T},y::Field{T},u, c,Wⁱ::Diagonal{T, Vector{T}},Qⁱ::Diagonal{T, Vector{T}},γ::Grid) where T <: Real
	squared model-data misfit for gridded data
	controls are a vector input for Optim.jl
	# Arguments
	- `convec`: concatenated control vecotr incuding u and f
	- `J`: cost function of sum of squared misfits
	- `gJ`: derivative of cost function wrt to controls
	- `u`: tracer controls, field format
	- `non_zero_indices`: Non-zero indices for reconstruction of water-mass matrix A
	- `b`: boundary conditions
	- `c`: tracer concentrations from GCM
	- `Wⁱ`: inverse of W weighting matrix for observations
	- `Qⁱ`: inverse of Q weighting matrix for tracer conservation
	- `γ`: grid
	"""
	function costfunction_gridded_model(convec,non_zero_indices,b₀::Union{BoundaryCondition{T},NamedTuple{<:Any, NTuple{N1,BoundaryCondition{T}}}},u₀::Field{T},y::Vector{T},c,Wⁱ::Diagonal{T, Vector{T}},Qⁱ::Diagonal{T, Vector{T}},γ::Grid) where {N1, N2, T <: Real}
	ulength = sum(γ.wet)
	uvec=convec[begin:ulength]
	non_zero_values = convec[ulength+1:end]
	A = sparse(non_zero_indices[:, 1], non_zero_indices[:, 2], non_zero_values)
	# find lagrange multipliers
	muk = transpose(A) * Qⁱ * (A * c)
	J = transpose(A * c) * Qⁱ * (A*c) + uvec ⋅ uvec - 2 * transpose(muk) * ((-Wⁱ * uvec - c + y)) #n ⋅ (Wⁱ * n) # dot product
	# adjoint equations
	guvec = zeros(length(convec))
	for (ii,vv) in enumerate(convec)
	if ii <= ulength
	#this is the derivative of the cost function wrt the part of the control vector
	# associated with the tracer concentration
	guvec[ii] = uvec[ii] + (2*transpose(muk) * Wⁱ)[ii]
	else
	#this is the derivative of the cost function wrt the part of the control vector
	# associated with the transport vector
	guvec[ii]=2 * Qⁱ[non_zero_indices[ii-ulength, 2],non_zero_indices[ii-ulength, 2]] * c[non_zero_indices[ii-ulength, 1]]^2 *convec[ii]
	end
	end
	return J , guvec
	end
	"""
	function costfunction_gridded_model!(J,guvec,convec::Vector{T},non_zero_indices,b₀::Union{BoundaryCondition{T},NamedTuple{<:Any, NTuple{N1,BoundaryCondition{T}}}},u₀::Union{BoundaryCondition{T},NamedTuple{<:Any, NTuple{N2,BoundaryCondition{T}}}},c,y::Field{T},Wⁱ::Diagonal{T, Vector{T}},Qⁱ::Diagonal{T, Vector{T}},γ::Grid) where {N1, N2, T <: Real}
	"""
	function costfunction_gridded_model!(J,guvec,convec::Vector{T},non_zero_indices,b₀::Union{BoundaryCondition{T},NamedTuple{<:Any, NTuple{N1,BoundaryCondition{T}}}},u₀::Field{T},y::Vector{T},c,Wⁱ::Diagonal{T, Vector{T}},Qⁱ::Diagonal{T, Vector{T}},γ::Grid) where {N1, N2, T <: Real}
	ulength = sum(γ.wet)
	uvec = convec[begin:ulength]
	non_zero_values = convec[ulength+1:end]
	A = sparse(non_zero_indices[:, 1], non_zero_indices[:, 2], non_zero_values)
	# find lagrange multipliers
	muk = transpose(A) * Qⁱ * (A * c)
	if guvec != nothing
	tmp = guvec
	for (ii,vv) in enumerate(tmp)
	if ii <= ulength
	guvec[ii] = uvec[ii] + (2 * transpose(muk) * Wⁱ)[ii]#vv
	else
	guvec[ii]=2 * Qⁱ[non_zero_indices[ii-ulength, 2],non_zero_indices[ii-ulength, 2]]* c[non_zero_indices[ii-ulength, 1]]^2 *convec[ii]
	end
	end
	end
	if J !=nothing
	return transpose(A * c) * Qⁱ * (A*c) + uvec ⋅ uvec -2 * transpose(muk) *(-Wⁱ * uvec - c + y )
	end
	end

(script to run this is here: https://github.com/ggebbie/TMI.jl/blob/invert-model/scripts/invert_model_TS.jl)

If I optimize right now the cost function grows, so I'm definitely doing something wrong. It was working ok until I implemented the Lagrange multipliers, so I probably have an error there. I'm mostly just posting this so that we can look at the code and talk about it at some point - see if I'm going in a good direction or if I need to refocus/rethink.

[ggebbie]

I'll take a look. I am currently working on issue #123 which has led me to refactor many things. I can help merge your work with the other changes.

[cspencerjones]

Updated list given that we are not learning the cross-face flux, but the water-mass matrix

• Replace cost function: tracer plus tracer control vector should equal model output of tracer (22 in the notes)
• Check that with noisy input we get back the model output of tracer
• Tracer control vector should be minimized (first term of 18)
• The vector f (and not the lu factorized version of A) should be an input to the cost function
• It then probably makes sense to add tracer conservation (17 in the notes) to the cost function
• Find the Lagrange multiplier for the tracer (26)
• A second control vector u_f should also be an input to the cost function
• Changes in u_f should be minimized (second term of 18)
• Add mass conservation to cost function (19) and to Lagrange multiplier
• Test to see if mass is conserved more accurately after optimization

[cspencerjones]

I ended up getting rid of some of the Lagrange multiplier code but I now think I'm optimizing successfully here: https://github.com/ggebbie/TMI.jl/tree/invert-model.

The next step is to try this on actual model output.