Python APIs for particles optimization (sampling and correspondence)

[sheryjoe]

@AtefehKashani Please define APIs and callbacks needed to perform particle optimization in python for the nonlinear shape model. Adding @cchriste to help Alan and Atefeh getting started with python APIs (pybind).

[sheryjoe]

BeforeIteration - compute shape statistics and correspondence object gradient. We can use python for training the network and compute the correspondence gradient. @AtefehKashani Which APIs are needed from the Optimize lib to train and compute gradients?

[AtefehKashani]

I have added a list of needed APIs here.

1. correspondence statistics, mean and covariance
2. gradient of correspondences in shape space
3. correspondences energy function to compute the amount of energy update at each iteration 4.sampling APIs
   • compute sampling energy
   • compute gradient (particle movement) based on sampling objective function
     • update particle based on a given gradient vector, which will entail projecting the update particle onto the surface.

[AtefehKashani]

Let's define some notations before going to the pseudocode.

• z: correspondences in each iteration, with dimension dM (where d: input dimension, and M: number of particles)
• z_0: mapping of z to prior space via f: z --> z_0, with dimension dM
• g : inverse of f function which maps back z_0 to z (g: z_0 --> z)
• NVP: invertable network
• Jac: Jacobian of g in respect to z_0

here is the steps in nonlinear shape modeling

• initialize z
• in each iteration do:
  • Compute the statistics of the shape space (z) (mean and covariance matrix)
  • Compute the gradient of the particles (z) (dH/dz)
    • update the prior space (z_0) mean and covariance matrix with the mean and covariance of the shape space (z)
    • learn NVP parameters using z of the current iteration and prior distribution
    • z_0 = NVP(z)
    • compute mean and covariance matrix z_0
    • dH/dz_0 : compute gradient of the particles in prior space (This can be done via SW API to compute the gradient)
    • dH/dz: compute the gradient of particles using Jac * dH/dz_0
  • Compute amount of Sampling update
  • Compute amount of particle update
    • map z to z_0 via learnt NVP for both old particle and new particles
    • H(z) = H(z_0) + mean(log |Jac|) for both old particle and new particles
    • H_{new}- H{old}
  • update particles
  • Update objective function
    • calculate sampling energy
    • calculate H(z)
      • map z to z_0 via learnt NVP
      • H(z) = H(z_0) + mean(log |Jac|)

[akenmorris]

@AtefehKashani, if you want to try the pybind_optimize branch, both methods of using the Optimize class with Python are implemented:

https://github.com/SCIInstitute/ShapeWorks/tree/pybind_optimize

Implemented so far:

• SetIterationCallbackFunction(func) set iteration callback (called on each iteration)
• GetParticleSystem() (returns a numpy array of the particle system)

You can drive from Python by doing this:

import numpy

import shapeworks
opt = shapeworks.Optimize()
opt.LoadParameterFile("optimize_python.xml")

def callback():
    print("python callback")
    particles = opt.GetParticleSystem()
    print(type(particles))
    print(particles)

opt.SetIterationCallbackFunction(callback)
opt.Run()

Alternatively, you can have ShapeWorksRun (and later Studio) drive by specifying:

<python_filename>file.py</python_filename>

This module will be loaded at startup and the "run" command will be called with the Optimize class as the parameter. An example file that accomplishes the same as above:

import shapeworks

opt = shapeworks.Optimize

def callback():
    print("python callback")
    particles = opt.GetParticleSystem()
    print(type(particles))
    print(particles)

def run(optimizer):
    global opt
    opt = optimizer
    opt.SetIterationCallbackFunction(callback)

Example output:

...
1. 0ms
python callback
<class 'numpy.ndarray'>
[[-19.12674606 -19.12674606 -23.51327091 -23.51327091]
 [-35.55056751 -35.55056751 -35.6341362  -35.6341362 ]
 [-91.48947597 -91.48947597 -90.42861462 -90.42861462]
 [-21.17568552 -21.17568552 -23.91829342 -23.91829342]
 [-39.20589983 -39.20589983 -37.9530549  -37.9530549 ]
 [-89.5236671  -89.5236671  -89.37270641 -89.37270641]]
Energy: 1.44793
...

---

Let me know what calls, both ways (to/from python) you would like added to the API.

pybind11 will automatically convert Eigen matrices to numpy.

[riddhishb]

Callback List

1. GetShapeWorksStage() → a flag that indicates whether shapeworks is in the initialize stage or the optimize stage. This is required as some of the processing is only in the optimize stage and need not be performed when particles are still in the initialize stage.
2. GetParticleMean(particles), GetParticleCovariance(particles) → generates the mean and covariance matrix for the given particle.
3. GetGradients(particles) → provides the gradient of the correspondence term (entropy) given a set of particles.
4. SetCorrespondenceGradients(gradients) → sets the gradient for the correspondence term.
5. SetCorrespondeceObjective(objective) → sets the objective function for the correspondence term.

[riddhishb]

I'll add in the callbacks here as I need them. I think we are good for now.

Callback	Status	Description
GetShapeWorksStage()	🕒	flag that indicates whether shapeworks is in the initialize stage or the optimize stage. This is required as some of the processing is only in the optimize stage and need not be performed when particles are still in the initialize stage
GetParticleSystem()	✅	Gets the current matrix containing the correspondence particles
GetCorrespondenceUpdateMatrix()	✅	Gets the correspondence updates for particles
GetParticles()	✅	Gets the current matrix containing the correspondence points
SetCorrespondenceUpdateMatrix(updateMatrix)	✅	Set the updates for the correspondence term from Python

[akenmorris]

I've added GetOptimizing() that returns true or false. True for optimization phase, false for initialization.

[riddhishb]

perfect!