optimizerDisplay.py

Go to the documentation of this file.

#
# LSST Data Management System
# Copyright 2008-2013 LSST Corporation.
#
# This product includes software developed by the
# LSST Project (http://www.lsst.org/).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the LSST License Statement and
# the GNU General Public License along with this program.  If not,
# see <http://www.lsstcorp.org/LegalNotices/>.
#
 
import numpy
import matplotlib
import matplotlib.colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
 
from .densityPlot import hide_xticklabels, hide_yticklabels
from .. import modelfitLib
 
__all__ = ("OptimizerDisplay", )
 
 
class OptimizerIterationDisplay:
 
    def __init__(self, parent, sample):
        self.parentparent = parent
        self.samplesample = sample
        # scale unit grid by trust radius
        self.gridgrid = parent.unitGrid * sample.get(parent.recorder.trust)
        # offset grid to center it on the current parameter point
        self.gridgrid += sample.get(parent.recorder.parameters).reshape((1,)*parent.ndim + (parent.ndim,))
        self._objectiveValues_objectiveValues = None
        self._objectiveModel_objectiveModel = None
        self.rejectedrejected = []
 
    def __getattr__(self, name):
        # look for keys on the recorder and lookup fields for unknown attributes
        return self.samplesample.get(getattr(self.parentparent.recorder, name))
 
    @property
    def objectiveValues(self):
        if self._objectiveValues_objectiveValues is None:
            self._objectiveValues_objectiveValues = numpy.zeros(self.gridgrid.shape[:-1], dtype=float)
            self.parentparent.objective.fillObjectiveValueGrid(self.gridgrid.reshape(-1, self.parentparent.ndim),
                                                         self._objectiveValues_objectiveValues.reshape(-1))
            good = numpy.isfinite(self._objectiveValues_objectiveValues)
            self._objectiveValues_objectiveValues[numpy.logical_not(good)] = self._objectiveValues_objectiveValues[good].max()
        return self._objectiveValues_objectiveValues
 
    @property
    def objectiveModel(self):
        if self._objectiveModel_objectiveModel is None:
            self._objectiveModel_objectiveModel = numpy.zeros(self.gridgrid.shape[:-1], dtype=float)
            self.parentparent.recorder.fillObjectiveModelGrid(self.samplesample,
                                                        self.gridgrid.reshape(-1, self.parentparent.ndim),
                                                        self._objectiveModel_objectiveModel.reshape(-1))
        return self._objectiveModel_objectiveModel
 
 
class OptimizerDisplay:
 
    def __init__(self, history, model, objective, steps=11):
        self.recorderrecorder = modelfitLib.OptimizerHistoryRecorder(history.schema)
        # len(dimensions) == N in comments below
        self.dimensionsdimensions = list(model.getNonlinearNames()) + list(model.getAmplitudeNames())
        self.ndimndim = len(self.dimensionsdimensions)
        self.tracktrack = []
        self.objectiveobjective = objective
        # This creates a array with shape [steps, ..., steps, N] (a total of N+1 array dimensions):
        # this is an N-dimensional grid, with the last dimension of the grid array giving the coordinates
        # of each grid point.
        # We slice mgrid to generate the basic grid, which is [N, steps, ..., steps]
        mgridArgs = (slice(-1.0, 1.0, steps*1j),) * self.ndimndim
        # We'll index the result of mgrid with these args to make first dimension last
        transposeArgs = tuple(list(range(1, self.ndimndim+1)) + [0])
        self.unitGridunitGrid = numpy.mgrid[mgridArgs].transpose(transposeArgs).copy()
        current = None
        for sample in history:
            if sample.get(self.recorderrecorder.state) & modelfitLib.Optimizer.STATUS_STEP_REJECTED:
                assert current is not None
                current.rejected.append(sample)
                continue
            current = OptimizerIterationDisplay(self, sample)
            self.tracktrack.append(current)
 
    def plot(self, xDim, yDim, n=0):
        return OptimizerDisplayFigure(self, xDim=xDim, yDim=yDim, n=n)
 
 
class OptimizerDisplayFigure:
 
    def __init__(self, parent, xDim, yDim, n=0):
        self.parentparent = parent
        self.xDimxDim = xDim
        self.yDimyDim = yDim
        self.jj = self.parentparent.dimensions.index(self.xDimxDim)
        self.ii = self.parentparent.dimensions.index(self.yDimyDim)
        self.yKeyyKey = self.parentparent.recorder.parameters[self.ii]
        self.xKeyxKey = self.parentparent.recorder.parameters[self.jj]
        self.zKeyzKey = self.parentparent.recorder.objective
        # grid slice indices corresponding to the dimensions we're plotting
        self.slice2dslice2d = [s//2 for s in self.parentparent.unitGrid.shape[:-1]]
        self.slice2dslice2d[self.ii] = slice(None)
        self.slice2dslice2d[self.jj] = slice(None)
        self.slice2dslice2d = tuple(self.slice2dslice2d)
        self.sliceXsliceX = [s//2 for s in self.parentparent.unitGrid.shape[:-1]]
        self.sliceXsliceX[self.jj] = slice(None)
        self.sliceXsliceX = tuple(self.sliceXsliceX)
        self.sliceYsliceY = [s//2 for s in self.parentparent.unitGrid.shape[:-1]]
        self.sliceYsliceY[self.ii] = slice(None)
        self.sliceYsliceY = tuple(self.sliceYsliceY)
        self.tracktrack = dict(
            x=numpy.array([iteration.sample.get(self.xKeyxKey) for iteration in self.parentparent.track]),
            y=numpy.array([iteration.sample.get(self.yKeyyKey) for iteration in self.parentparent.track]),
            z=numpy.array([iteration.sample.get(self.zKeyzKey) for iteration in self.parentparent.track]),
        )
        self.nn = n
        self.figurefigure = matplotlib.pyplot.figure(f"{xDim} vs {yDim}", figsize=(16, 8))
        self.figurefigure.subplots_adjust(left=0.025, right=0.975, bottom=0.08, top=0.95, wspace=0.12)
        self.axes3daxes3d = self.figurefigure.add_subplot(1, 2, 1, projection='3d')
        self.axes3daxes3d.autoscale(False)
        self.axes3daxes3d.set_xlabel(self.xDimxDim)
        self.axes3daxes3d.set_ylabel(self.yDimyDim)
        self.axes2daxes2d = self.figurefigure.add_subplot(1, 2, 2)
        self.axes2daxes2d.set_xlabel(self.xDimxDim)
        self.axes2daxes2d.set_ylabel(self.yDimyDim)
        self.axes2daxes2d.autoscale(False)
        divider = make_axes_locatable(self.axes2daxes2d)
        self.axesXaxesX = divider.append_axes("top", 1.5, pad=0.1, sharex=self.axes2daxes2d)
        self.axesXaxesX.autoscale(False, axis='x')
        hide_xticklabels(self.axesXaxesX)
        self.axesYaxesY = divider.append_axes("right", 1.5, pad=0.1, sharey=self.axes2daxes2d)
        self.axesYaxesY.autoscale(False, axis='y')
        hide_yticklabels(self.axesYaxesY)
        self.artistsartists = []
        self.guessExtentguessExtent()
        self.plotTrackplotTrack()
        self.plotRejectedplotRejected()
        self.plotSurfacesplotSurfaces()
 
    @property
    def xlim(self):
        return self._extent_extent[:2]
 
    @property
    def ylim(self):
        return self._extent_extent[2:4]
 
    @property
    def zlim(self):
        return self._extent_extent[4:]
 
    def guessExtent(self):
        current = self.parentparent.track[self.nn]
        x = current.sample.get(self.xKeyxKey)
        y = current.sample.get(self.yKeyyKey)
        zMin1 = current.objectiveValues[self.slice2dslice2d].min()
        zMax1 = current.objectiveValues[self.slice2dslice2d].max()
        zMin2 = current.objectiveModel[self.slice2dslice2d].min()
        zMax2 = current.objectiveModel[self.slice2dslice2d].max()
        self.setExtentsetExtent(x0=x - current.trust, x1=x + current.trust,
                       y0=y - current.trust, y1=y + current.trust,
                       z0=min(zMin1, zMin2), z1=max(zMax1, zMax2), lock=False)
 
    def setExtent(self, x0=None, x1=None, y0=None, y1=None, z0=None, z1=None, lock=True):
        if x0 is None:
            x0 = self._extent_extent[0]
        if x1 is None:
            x1 = self._extent_extent[1]
        if y0 is None:
            y0 = self._extent_extent[2]
        if y1 is None:
            y1 = self._extent_extent[3]
        if z0 is None:
            z0 = self._extent_extent[4]
        if z1 is None:
            z1 = self._extent_extent[5]
        self._extent_extent = (x0, x1, y0, y1, z0, z1)
        self._lock_lock = lock
        self.axes3daxes3d.set_xlim(*self.xlimxlim)
        self.axes3daxes3d.set_ylim(*self.ylimylim)
        self.axes3daxes3d.set_zlim(*self.zlimzlim)
        self.axes2daxes2d.set_xlim(*self.xlimxlim)
        self.axes2daxes2d.set_ylim(*self.ylimylim)
        self.axesXaxesX.set_ylim(*self.zlimzlim)
        self.axesYaxesY.set_xlim(*self.zlimzlim)
 
    def _clipZ(self, x, y, z):
        # clipping is currently disabled; more trouble than it's worth
        if False:
            mask = numpy.logical_or.reduce([x < self.xlimxlim[0], x > self.xlimxlim[1],
                                            y < self.ylimylim[0], y > self.ylimylim[1],
                                            z < self.zlimzlim[0], z > self.zlimzlim[1]],
                                           axis=0)
 
            z[mask] = numpy.nan
            return numpy.logical_not(mask).astype(int).sum() > 4
        return True
 
    def _contour(self, axes, *args, **kwds):
        self.artistsartists.extend(axes.contour(*args, **kwds).collections)
 
    def plotTrack(self):
        kwds = dict(markeredgewidth=0, markerfacecolor='g', color='g', marker='o')
        self.axes3daxes3d.plot(self.tracktrack['x'], self.tracktrack['y'], self.tracktrack['z'], **kwds)
        self.axes2daxes2d.plot(self.tracktrack['x'], self.tracktrack['y'], **kwds)
        self.axesXaxesX.plot(self.tracktrack['x'], self.tracktrack['z'], **kwds)
        self.axesYaxesY.plot(self.tracktrack['z'], self.tracktrack['y'], **kwds)
 
    def plotRejected(self):
        kwds = dict(markeredgewidth=0, markerfacecolor='r', color='r', marker='v')
        current = self.parentparent.track[self.nn]
        cx = current.sample.get(self.xKeyxKey)
        cy = current.sample.get(self.yKeyyKey)
        cz = current.sample.get(self.zKeyzKey)
        for r in current.rejected:
            x = [cx, r.get(self.xKeyxKey)]
            y = [cy, r.get(self.yKeyyKey)]
            z = [cz, r.get(self.zKeyzKey)]
            self.artistsartists.extend(self.axes3daxes3d.plot(x, y, z, **kwds))
            self.artistsartists.extend(self.axes2daxes2d.plot(x, y, **kwds))
            self.artistsartists.extend(self.axesXaxesX.plot(x, z, **kwds))
            self.artistsartists.extend(self.axesYaxesY.plot(z, y, **kwds))
 
    def plotSurfaces(self):
        current = self.parentparent.track[self.nn]
 
        # Start with 2-d and 3-d surfaces
        x = current.grid[self.slice2dslice2d + (self.jj,)]
        y = current.grid[self.slice2dslice2d + (self.ii,)]
        z1 = current.objectiveValues[self.slice2dslice2d].copy()
        z2 = current.objectiveModel[self.slice2dslice2d].copy()
        norm = matplotlib.colors.Normalize(vmin=self.zlimzlim[0], vmax=self.zlimzlim[1])
 
        self._contour_contour(self.axes2daxes2d, x, y, z1, cmap=matplotlib.cm.spring, norm=norm)
        self._contour_contour(self.axes2daxes2d, x, y, z2, cmap=matplotlib.cm.winter, norm=norm)
 
        # matplotlib doesn't do clipping in 3d, so we'll do that manually
        if self._clipZ_clipZ(x, y, z1):
            self._contour_contour(self.axes3daxes3d, x, y, z1, cmap=matplotlib.cm.spring, norm=norm)
            self.artistsartists.append(self.axes3daxes3d.plot_surface(x, y, z1, rstride=1, cstride=1,
                                                         cmap=matplotlib.cm.spring, norm=norm,
                                                         linewidth=0, antialiased=1, alpha=0.5))
        if self._clipZ_clipZ(x, y, z2):
            self._contour_contour(self.axes3daxes3d, x, y, z2, cmap=matplotlib.cm.winter, norm=norm)
            self.artistsartists.append(self.axes3daxes3d.plot_surface(x, y, z2, rstride=1, cstride=1,
                                                         cmap=matplotlib.cm.winter, norm=norm,
                                                         linewidth=0, antialiased=1, alpha=0.5))
 
        # Now the 1-d surfaces
        self.artistsartists.extend(self.axesXaxesX.plot(current.grid[self.sliceXsliceX + (self.jj,)],
                                            current.objectiveValues[self.sliceXsliceX], 'm-'))
        self.artistsartists.extend(self.axesXaxesX.plot(current.grid[self.sliceXsliceX + (self.jj,)],
                                            current.objectiveModel[self.sliceXsliceX], 'c-'))
        self.artistsartists.extend(self.axesYaxesY.plot(current.objectiveValues[self.sliceYsliceY],
                                            current.grid[self.sliceYsliceY + (self.ii,)], 'm-'))
        self.artistsartists.extend(self.axesYaxesY.plot(current.objectiveModel[self.sliceYsliceY],
                                            current.grid[self.sliceYsliceY + (self.ii,)], 'c-'))
 
    def move(self, n):
        self.nn = n
        if not self._lock_lock:
            self.guessExtentguessExtent()
        for artist in self.artistsartists:
            try:
                artist.remove()
            except TypeError:
                # sometimes matplotlib throws an exception even though everything worked fine
                pass
        self.artistsartists = []
        self.plotSurfacesplotSurfaces()
        self.plotRejectedplotRejected()
        self.figurefigure.canvas.draw()