check_partials not working as intended with sparse partials

[johnjasa]

Description

Originally sent to me as an email from @cfe316, I'm opening an issue so this can be tracked and solved. From Jacob:

"""
The docs say that OM will pick up on the sparsity pattern, but I'm worried there's a problem with this specifically during the check_partials step.

Evidence

1. I'm declaring partials with self.declare_partials("P", "xu", val=<sparse matrix>). I see that when I run check_partials on it,
   1a) J_fd is missing certain entries. J_fwd is correct.
   1b) The handy plot_partials function breaks.
   1c) Weirdly when I run check_partials a second time, the J_fwd matrix that gets printed also drops those entries.

2. I can fix this by using self.declare_partials("P", "xu", rows=r, cols=c, val=data) where I've previously extracted r, c, and data from the sparse matrix.

Q1) Does the sparsity pattern pickup only work some formats of sparse matrix? i.e. CSR but not COO?
Q2) Are the docs a lie???? [🎂 cake emoji added by John]
"""

I've poked around a bit in problem.py where we define check_partials and I haven't found anything conclusive. It smells like some of the logic for how the deriv_value is calculated (around here or before) might be relevant?

Jacob very helpfully provided a MWE with comments to flip on or off the behavior.

Example

import openmdao.api as om
import numpy as np
from scipy.sparse import lil_matrix

def quadpoly_shapecomp(shapes):
    xu_length = shapes['xu'][0]
    quads_length = xu_length - 1
    return (quads_length, 4, 2)

class QuadPoly(om.ExplicitComponent):

    def setup(self):
        self.add_input('xu', shape_by_conn=True)
        self.add_input('yu', shape_by_conn=True, copy_shape='xu')        
        self.add_input('xl', shape_by_conn=True, copy_shape='xu')
        self.add_input('yl', shape_by_conn=True, copy_shape='xu')

        self.add_output('P', compute_shape=quadpoly_shapecomp)

    def setup_partials(self):
        length_x = self._get_var_meta('xu', 'size')
        length_p = length_x - 1
        flattened_shape = (length_p * 4 * 2, length_x)
            
        dP_dxl = lil_matrix((8 * length_p, length_x), dtype=int)
        dP_dyl = lil_matrix((8 * length_p, length_x), dtype=int)
        for i in range(length_p):
            dP_dxl[8*i + 0, i] = 1
            dP_dxl[8*i + 2, i + 1] = 1
            dP_dyl[8*i + 1, i] = 1
            dP_dyl[8*i + 3, i + 1] = 1

        # This doesn't work: 2nd column is all zeros in J_fd below!
        # Also this makes the plot_partial_derivs fail for some reason
        self.declare_partials('P', 'xl', val=dP_dxl.tocsr())

        # Here is a way to make it work by explicitly providing rows, cols as well.
        # dP_dxl_csr = dP_dxl.tocsr()
        # r, c = dP_dxl_csr.nonzero()
        # data = dP_dxl_csr.data
        # self.declare_partials('P', 'xl', rows=r, cols=c, val=data)

        # this does work for some reason, even though it looks the same to me.
        dP_dyl_csr = dP_dyl.tocsr()
        self.declare_partials('P', 'yl', val=dP_dyl_csr)
        
        # This works fine
        mtxu = np.zeros((length_p, 4, 2, length_x))
        mtyu = np.zeros((length_p, 4, 2, length_x))
        for i in range(length_p):
            mtxu[i, 3, 0, i] = 1 
            mtxu[i, 2, 0, i+1] = 1
            mtyu[i, 3, 1, i] = 1
            mtyu[i, 2, 1, i+1] = 1

        self.declare_partials('P', 'xu', val=mtxu.reshape(flattened_shape))
        self.declare_partials('P', 'yu', val=mtyu.reshape(flattened_shape))
    
    def compute(self, inputs, outputs):
        xu = inputs['xu']
        xl = inputs['xl']
        yu = inputs['yu']
        yl = inputs['yl']

        def lefts_and_rights(x):
            return x[:-1], x[1:]

        length_x = self._get_var_meta('xu', 'size')
        length_p = length_x - 1
            
        xl0, xl1 = lefts_and_rights(xl)
        xu0, xu1 = lefts_and_rights(xu)
        yl0, yl1 = lefts_and_rights(yl)
        yu0, yu1 = lefts_and_rights(yu)
        
        poly = np.zeros((length_p, 4, 2))
        poly[:, 0, 0] = xl0
        poly[:, 1, 0] = xl1
        poly[:, 2, 0] = xu1
        poly[:, 3, 0] = xu0
        poly[:, 0, 1] = yl0
        poly[:, 1, 1] = yl1
        poly[:, 2, 1] = yu1
        poly[:, 3, 1] = yu0
        
        outputs['P'] = poly


class MyGroup2(om.Group):

    def setup(self):
        self.add_subsystem('quadpoly', QuadPoly())

prob = om.Problem()
model = prob.model

indeps = model.add_subsystem('indeps', om.IndepVarComp())
indeps.add_output('xl', [0.0, 1.23, 2.0])
indeps.add_output('xu', [0.0, 1.0, 2.0])
indeps.add_output('yl', [-1.0, -1.0, -1.0,])
indeps.add_output('yu', [1.0, 1.0, 1.0,])

prob.model.add_subsystem('grp', MyGroup2(), promotes=['*'])

model.connect('indeps.xl', 'quadpoly.xl')
model.connect('indeps.xu', 'quadpoly.xu')
model.connect('indeps.yl', 'quadpoly.yl')
model.connect('indeps.yu', 'quadpoly.yu')

prob.setup()
prob.run_model()

check_partials_data = prob.check_partials(compact_print=True, out_stream=None)
print(check_partials_data['grp.quadpoly'][('P', 'xl')])
# check_partials looks bad: try uncommenting the relevant part in 

OpenMDAO Version

3.31.2-dev

Relevant environment information

No response

[naylor-b]

The easiest way to fix this is to just convert the sparse array to COO when they call declare_partials, and store rows, cols, and val in our internal coo format. Would that be acceptable @johnjasa @robfalck @cfe316 ?

[robfalck]

I think that's the right way to do it. Check if Val is a sparse array and convert as necessary, otherwise use rows and cols and assume it's COO

[cfe316]

Hi Rob, John, Bret,

Thanks for your advice. Right now I'm using the more 'robust' method of manually getting the nonzero entries and specifying rows=r, cols=c, val=data. This seems to be fine for the moment, but it would be nice to just pass the sparse array into val=<sparse array> directly.

(I don't have any opinion on which flavor of sparse array I should be using.)