text:
#include <petsc/private/pcimpl.h>     /*I "petscpc.h" I*/
#include <petsc.h>
#include <petsc/private/hashmapi.h>
#include <petscsf.h>
#include <libssc.h>

PetscLogEvent PC_Patch_CreatePatches, PC_Patch_ComputeOp, PC_Patch_Solve, PC_Patch_Scatter, PC_Patch_Apply, PC_Patch_Prealloc;

static PetscBool PCPatchPackageInitialized = PETSC_FALSE;

PETSC_EXTERN PetscErrorCode PCPatchInitializePackage(void)
{
    PetscErrorCode ierr;
    PetscFunctionBegin;

    if (PCPatchPackageInitialized) PetscFunctionReturn(0);
    PCPatchPackageInitialized = PETSC_TRUE;
    ierr = PCRegister("patch", PCCreate_PATCH); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHCreate", PC_CLASSID, &PC_Patch_CreatePatches); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHComputeOp", PC_CLASSID, &PC_Patch_ComputeOp); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHSolve", PC_CLASSID, &PC_Patch_Solve); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHApply", PC_CLASSID, &PC_Patch_Apply); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHScatter", PC_CLASSID, &PC_Patch_Scatter); CHKERRQ(ierr);
    ierr = PetscLogEventRegister("PCPATCHPrealloc", PC_CLASSID, &PC_Patch_Prealloc); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

typedef struct {
    PetscSF          defaultSF;
    PetscSection    *dofSection;
    PetscSection     cellCounts;
    PetscSection     cellNumbering; /* Numbering of cells in DM */
    PetscSection     gtolCounts;   /* Indices to extract from local to
                                   * patch vectors */
    PetscInt         nsubspaces;   /* for mixed problems */
    PetscInt        *subspaceOffsets; /* offsets for calculating concatenated numbering for mixed spaces */
    PetscSection     bcCounts;
    IS               cells;
    IS               dofs;
    IS               ghostBcNodes;
    IS               globalBcNodes;
    IS               gtol;

    PetscBool        save_operators; /* Save all operators (or create/destroy one at a time?) */
    PetscBool        partition_of_unity; /* Weight updates by dof multiplicity? */
    PetscBool        multiplicative; /* Gauss-Seidel or Jacobi? */
    PetscInt         npatch;     /* Number of patches */
    PetscInt        *bs;            /* block size (can come from global
                                    * operators?) */
    PetscInt        *nodesPerCell;
    PetscInt         totalDofsPerCell;
    const PetscInt **cellNodeMap; /* Map from cells to nodes */

    KSP             *ksp;        /* Solvers for each patch */
    Vec              localX, localY;
    Vec              dof_weights; /* In how many patches does each dof lie? */
    Vec             *patchX, *patchY; /* Work vectors for patches */
    Vec             *patch_dof_weights;
    Mat             *mat;        /* Operators */
    MatType          sub_mat_type;
    PetscErrorCode  (*usercomputeop)(PC, Mat, PetscInt, const PetscInt *, PetscInt, const PetscInt *, void *);
    void            *usercomputectx;

    PetscErrorCode  (*patchconstructop)(void*, DM, PetscInt, PetscHMapI); /* patch construction */
    PetscInt         codim; /* dimension or codimension of entities to loop over; */
    PetscInt         dim;   /* only one of them can be set */
    PetscInt         exclude_subspace; /* If you don't want any other dofs from a particular subspace you can exclude them with this.
                                          Used for Vanka in Stokes, for example, to eliminate all pressure dofs not on the vertex
                                          you're building the patch around */
    PetscInt         vankadim;   /* In Vanka construction, should we eliminate any entities of a certain dimension? */

    PetscBool        print_patches; /* Should we print out information about patch construction? */
    PetscBool        symmetrise_sweep; /* Should we sweep forwards->backwards, backwards->forwards? */

    IS              *userIS;
    IS               iterationSet; /* Index set specifying how we iterate over patches */
    PetscInt         nuserIS; /* user-specified index sets to specify the patches */
    PetscBool        user_patches;
    PetscErrorCode  (*userpatchconstructionop)(PC, PetscInt*, IS**, IS*, void* ctx);
    void            *userpatchconstructctx;
} PC_PATCH;

PETSC_EXTERN PetscErrorCode PCPatchSetSaveOperators(PC pc, PetscBool flg)
{
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscFunctionBegin;

    patch->save_operators = flg;
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetPartitionOfUnity(PC pc, PetscBool flg)
{
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscFunctionBegin;

    patch->partition_of_unity = flg;
    PetscFunctionReturn(0);
}

static PetscErrorCode PCPatchCreateDefaultSF_Private(PC pc, PetscInt n, const PetscSF *sf, const PetscInt *bs)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscFunctionBegin;

    if (n == 1 && bs[0] == 1) {
        patch->defaultSF = sf[0];
        ierr = PetscObjectReference((PetscObject)patch->defaultSF); CHKERRQ(ierr);
    } else {
        PetscInt allRoots = 0, allLeaves = 0;
        PetscInt leafOffset = 0;
        PetscInt *ilocal = NULL;
        PetscSFNode *iremote = NULL;
        PetscInt *remoteOffsets = NULL;
        PetscInt index = 0;
        PetscHMapI rankToIndex;
        PetscInt numRanks = 0;
        PetscSFNode *remote = NULL;
        PetscSF rankSF;
        PetscInt *ranks = NULL;
        PetscInt *offsets = NULL;
        MPI_Datatype contig;
        PetscHMapI ht;

        /* First figure out how many dofs there are in the concatenated numbering.
         * allRoots: number of owned global dofs;
         * allLeaves: number of visible dofs (global + ghosted).
         */
        for ( PetscInt i = 0; i < n; i++ ) {
            PetscInt nroots, nleaves;
            ierr = PetscSFGetGraph(sf[i], &nroots, &nleaves, NULL, NULL); CHKERRQ(ierr);
            allRoots += nroots * bs[i];
            allLeaves += nleaves * bs[i];
        }
        ierr = PetscMalloc1(allLeaves, &ilocal); CHKERRQ(ierr);
        ierr = PetscMalloc1(allLeaves, &iremote); CHKERRQ(ierr);

        /* Now build an SF that just contains process connectivity. */
        PetscHMapICreate(&ht);
        for (PetscInt i = 0; i < n; i++ ) {
            PetscInt nranks;
            const PetscMPIInt *ranks = NULL;
            ierr = PetscSFSetUp(sf[i]); CHKERRQ(ierr);
            ierr = PetscSFGetRanks(sf[i], &nranks, &ranks, NULL, NULL, NULL); CHKERRQ(ierr);
            /* These are all the ranks who communicate with me. */
            for (PetscInt j = 0; j < nranks; j++) {
                PetscHMapISet(ht, (PetscInt)ranks[j], 0);
            }
        }
        PetscHMapIGetSize(ht, &numRanks); CHKERRQ(ierr);
        ierr = PetscMalloc1(numRanks, &remote); CHKERRQ(ierr);
        ierr = PetscMalloc1(numRanks, &ranks); CHKERRQ(ierr);
        ierr = PetscHMapIGetKeys(ht, &index, ranks); CHKERRQ(ierr);

        PetscHMapICreate(&rankToIndex);
        for (PetscInt i = 0; i < numRanks; i++) {
            remote[i].rank = ranks[i];
            remote[i].index = 0;
            PetscHMapISet(rankToIndex, ranks[i], i);
        }
        ierr = PetscFree(ranks); CHKERRQ(ierr);
        PetscHMapIDestroy(&ht);
        ierr = PetscSFCreate(PetscObjectComm((PetscObject)pc), &rankSF); CHKERRQ(ierr);
        ierr = PetscSFSetGraph(rankSF, 1, numRanks, NULL, PETSC_OWN_POINTER, remote, PETSC_OWN_POINTER); CHKERRQ(ierr);
        ierr = PetscSFSetUp(rankSF); CHKERRQ(ierr);

        /* OK, use it to communicate the root offset on the remote
         * processes for each subspace. */
        ierr = PetscMalloc1(n, &offsets); CHKERRQ(ierr);
        ierr = PetscMalloc1(n*numRanks, &remoteOffsets); CHKERRQ(ierr);

        offsets[0] = 0;
        for (PetscInt i = 1; i < n; i++) {
            PetscInt nroots;
            ierr = PetscSFGetGraph(sf[i-1], &nroots, NULL, NULL, NULL); CHKERRQ(ierr);
            offsets[i] = offsets[i-1] + nroots*bs[i-1];
        }
        /* Offsets are the offsets on the current process of the
         * global dof numbering for the subspaces. */
        ierr = MPI_Type_contiguous(n, MPIU_INT, &contig); CHKERRQ(ierr);
        ierr = MPI_Type_commit(&contig); CHKERRQ(ierr);

        ierr = PetscSFBcastBegin(rankSF, contig, offsets, remoteOffsets); CHKERRQ(ierr);
        ierr = PetscSFBcastEnd(rankSF, contig, offsets, remoteOffsets); CHKERRQ(ierr);
        ierr = MPI_Type_free(&contig); CHKERRQ(ierr);
        ierr = PetscFree(offsets); CHKERRQ(ierr);
        ierr = PetscSFDestroy(&rankSF); CHKERRQ(ierr);
        /* Now remoteOffsets contains the offsets on the remote
         * processes who communicate with me.  So now we can
         * concatenate the list of SFs into a single one. */
        index = 0;
        for ( PetscInt i = 0; i < n; i++ ) {
            PetscInt nroots, nleaves;
            const PetscInt *local = NULL;
            const PetscSFNode *remote = NULL;
            ierr = PetscSFGetGraph(sf[i], &nroots, &nleaves, &local, &remote); CHKERRQ(ierr);
            for ( PetscInt j = 0; j < nleaves; j++ ) {
                PetscInt rank = remote[j].rank;
                PetscInt idx, rootOffset;
                PetscHMapIGet(rankToIndex, rank, &idx);
                if (idx == -1) {
                    SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Didn't find rank, huh?");
                }
                /* Offset on given rank for ith subspace */
                rootOffset = remoteOffsets[n*idx + i];
                for ( PetscInt k = 0; k < bs[i]; k++ ) {
                    ilocal[index] = (local ? local[j] : j)*bs[i] + k + leafOffset;
                    iremote[index].rank = remote[j].rank;
                    iremote[index].index = remote[j].index*bs[i] + k + rootOffset;
                    ++index;
                }
            }
            leafOffset += nleaves * bs[i];
        }
        PetscHMapIDestroy(&rankToIndex);
        ierr = PetscFree(remoteOffsets); CHKERRQ(ierr);
        ierr = PetscSFCreate(PetscObjectComm((PetscObject)pc), &patch->defaultSF); CHKERRQ(ierr);
        ierr = PetscSFSetGraph(patch->defaultSF, allRoots, allLeaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER); CHKERRQ(ierr);
    }
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetCellNumbering(PC pc, PetscSection cellNumbering)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscFunctionBegin;

    patch->cellNumbering = cellNumbering;
    ierr = PetscObjectReference((PetscObject)cellNumbering); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetDiscretisationInfo(PC pc, PetscInt nsubspaces,
                                                         DM *dms,
                                                         PetscInt *bs,
                                                         PetscInt *nodesPerCell,
                                                         const PetscInt **cellNodeMap,
                                                         const PetscInt *subspaceOffsets,
                                                         PetscInt numGhostBcs,
                                                         const PetscInt *ghostBcNodes,
                                                         PetscInt numGlobalBcs,
                                                         const PetscInt *globalBcNodes)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscSF        *sfs;
    PetscFunctionBegin;

    ierr = PetscMalloc1(nsubspaces, &sfs); CHKERRQ(ierr);
    ierr = PetscMalloc1(nsubspaces, &patch->dofSection); CHKERRQ(ierr);
    ierr = PetscMalloc1(nsubspaces, &patch->bs); CHKERRQ(ierr);
    ierr = PetscMalloc1(nsubspaces, &patch->nodesPerCell); CHKERRQ(ierr);
    ierr = PetscMalloc1(nsubspaces, &patch->cellNodeMap); CHKERRQ(ierr);
    ierr = PetscMalloc1(nsubspaces+1, &patch->subspaceOffsets); CHKERRQ(ierr);

    patch->nsubspaces = nsubspaces;
    patch->totalDofsPerCell = 0;
    for (int i = 0; i < nsubspaces; i++) {
        ierr = DMGetDefaultSection(dms[i], &patch->dofSection[i]); CHKERRQ(ierr);
        ierr = PetscObjectReference((PetscObject)patch->dofSection[i]); CHKERRQ(ierr);
        patch->bs[i] = bs[i];
        patch->nodesPerCell[i] = nodesPerCell[i];
        patch->totalDofsPerCell += nodesPerCell[i]*bs[i];
        patch->cellNodeMap[i] = cellNodeMap[i];
        patch->subspaceOffsets[i] = subspaceOffsets[i];
        ierr = DMGetDefaultSF(dms[i], &sfs[i]); CHKERRQ(ierr);
    }
    ierr = PCPatchCreateDefaultSF_Private(pc, nsubspaces, sfs, patch->bs); CHKERRQ(ierr);
    ierr = PetscFree(sfs); CHKERRQ(ierr);

    patch->subspaceOffsets[nsubspaces] = subspaceOffsets[nsubspaces];
    ierr = ISCreateGeneral(PETSC_COMM_SELF, numGhostBcs, ghostBcNodes, PETSC_COPY_VALUES, &patch->ghostBcNodes); CHKERRQ(ierr);
    ierr = ISCreateGeneral(PETSC_COMM_SELF, numGlobalBcs, globalBcNodes, PETSC_COPY_VALUES, &patch->globalBcNodes); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetSubMatType(PC pc, MatType sub_mat_type)
{
    PetscErrorCode ierr;
    PC_PATCH      *patch = (PC_PATCH *)pc->data;
    PetscFunctionBegin;
    if (patch->sub_mat_type) {
        ierr = PetscFree(patch->sub_mat_type); CHKERRQ(ierr);
    }
    ierr = PetscStrallocpy(sub_mat_type, (char **)&patch->sub_mat_type); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetComputeOperator(PC pc, PetscErrorCode (*func)(PC, Mat, PetscInt,
                                                                                const PetscInt *,
                                                                                PetscInt,
                                                                                const PetscInt *,
                                                                                void *),
                                                      void *ctx)
{
    PC_PATCH *patch = (PC_PATCH *)pc->data;

    PetscFunctionBegin;
    /* User op can assume matrix is zeroed */
    patch->usercomputeop = func;
    patch->usercomputectx = ctx;

    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchSetUserPatchConstructionOperator(PC pc, PetscErrorCode (*func)(PC, PetscInt*, IS**, IS*, void*), void* ctx)
{
    PC_PATCH *patch = (PC_PATCH *)pc->data;

    PetscFunctionBegin;
    patch->userpatchconstructionop = func;
    patch->userpatchconstructctx = ctx;

    PetscFunctionReturn(0);
}

/* On entry, ht contains the topological entities whose dofs we are responsible for solving for;
   on exit, cht contains all the topological entities we need to compute their residuals.
   In full generality this should incorporate knowledge of the sparsity pattern of the matrix;
   here we assume a standard FE sparsity pattern.*/
static PetscErrorCode PCPatchCompleteCellPatch(DM dm, PetscHMapI ht, PetscHMapI cht)
{
    PetscErrorCode    ierr;
    PetscHashIter    hi;
    PetscInt          entity;
    PetscInt         *star = NULL, *closure = NULL;

    PetscFunctionBegin;

    PetscHMapIClear(cht);
    PetscHashIterBegin(ht, hi);
    while (!PetscHashIterAtEnd(ht, hi)) {
        PetscInt       starSize, closureSize;

        PetscHashIterGetKey(ht, hi, entity);
        PetscHashIterNext(ht, hi);

        /* Loop over all the cells that this entity connects to */
        ierr = DMPlexGetTransitiveClosure(dm, entity, PETSC_FALSE, &starSize, &star); CHKERRQ(ierr);
        for ( PetscInt si = 0; si < starSize; si++ ) {
            PetscInt ownedentity = star[2*si];
            /* now loop over all entities in the closure of that cell */
            ierr = DMPlexGetTransitiveClosure(dm, ownedentity, PETSC_TRUE, &closureSize, &closure); CHKERRQ(ierr);
            for ( PetscInt ci = 0; ci < closureSize; ci++ ) {
                PetscInt seenentity = closure[2*ci];
                PetscHMapISet(cht, seenentity, 0);
            }
        }
    }
    /* Only restore work arrays at very end. */
    if (closure) {
        ierr = DMPlexRestoreTransitiveClosure(dm, 0, PETSC_TRUE, NULL, &closure); CHKERRQ(ierr);
    }
    if (star) {
        ierr = DMPlexRestoreTransitiveClosure(dm, 0, PETSC_FALSE, NULL, &star); CHKERRQ(ierr);
    }
    PetscFunctionReturn(0);
}

/* Given a hash table with a set of topological entities (pts), compute the degrees of
   freedom in global concatenated numbering on those entities.
   For Vanka smoothing, this needs to do something special: ignore dofs of the
   constraint subspace on entities that aren't the base entity we're building the patch
   around. */
static PetscErrorCode PCPatchGetPointDofs(PC_PATCH *patch, PetscHMapI pts, PetscHMapI dofs, PetscInt base, PetscInt exclude_subspace)
{
    PetscErrorCode    ierr;
    PetscInt          ldof, loff;
    PetscHashIter    hi;
    PetscInt          p;

    PetscFunctionBegin;
    PetscHMapIClear(dofs);

    for ( PetscInt k = 0; k < patch->nsubspaces; k++ ) {
        PetscSection dofSection = patch->dofSection[k];
        PetscInt bs = patch->bs[k];
        PetscInt subspaceOffset = patch->subspaceOffsets[k];

        if (k == exclude_subspace) {
            /* only get this subspace dofs at the base entity, not any others */
            ierr = PetscSectionGetDof(dofSection, base, &ldof); CHKERRQ(ierr);
            ierr = PetscSectionGetOffset(dofSection, base, &loff); CHKERRQ(ierr);
            if (0 == ldof) continue;
            for ( PetscInt j = loff; j < ldof + loff; j++ ) {
                for ( PetscInt l = 0; l < bs; l++ ) {
                    PetscInt dof = bs*j + l + subspaceOffset;
                    PetscHMapISet(dofs, dof, 0);
                }
            }
            continue; /* skip the other dofs of this subspace */
        }

        PetscHashIterBegin(pts, hi);
        while (!PetscHashIterAtEnd(pts, hi)) {
            PetscHashIterGetKey(pts, hi, p);
            PetscHashIterNext(pts, hi);
            ierr = PetscSectionGetDof(dofSection, p, &ldof); CHKERRQ(ierr);
            ierr = PetscSectionGetOffset(dofSection, p, &loff); CHKERRQ(ierr);
            if (0 == ldof) continue;
            for ( PetscInt j = loff; j < ldof + loff; j++ ) {
                for ( PetscInt l = 0; l < bs; l++ ) {
                    PetscInt dof = bs*j + l + subspaceOffset;
                    PetscHMapISet(dofs, dof, 0);
                }
            }
        }
    }

    PetscFunctionReturn(0);
}

/* Given two hash tables A and B, compute the keys in B that are not in A, and
   put them in C */
static PetscErrorCode PCPatchComputeSetDifference(PetscHMapI A, PetscHMapI B, PetscHMapI C)
{
    PetscHashIter    hi;
    PetscInt          key;
    PetscBool         flg;

    PetscFunctionBegin;
    PetscHMapIClear(C);

    PetscHashIterBegin(B, hi);
    while (!PetscHashIterAtEnd(B, hi)) {
        PetscHashIterGetKey(B, hi, key);
        PetscHashIterNext(B, hi);
        PetscHMapIHas(A, key, &flg);
        if (!flg) {
            PetscHMapISet(C, key, 0);
        }
    }

    PetscFunctionReturn(0);
}

/*
 * PCPatchCreateCellPatches - create patches.
 *
 * Input Parameters:
 * + dm - The DMPlex object defining the mesh
 *
 * Output Parameters:
 * + cellCounts - Section with counts of cells around each vertex
 * - cells - IS of the cell point indices of cells in each patch
 */
static PetscErrorCode PCPatchCreateCellPatches(PC pc)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch      = (PC_PATCH *)pc->data;
    DM              dm;
    DMLabel         ghost;
    PetscInt        pStart, pEnd, vStart, vEnd, cStart, cEnd;
    PetscBool       flg;
    PetscInt       *cellsArray = NULL;
    PetscInt        numCells;
    PetscSection    cellCounts;
    PetscHMapI      ht;
    PetscHMapI      cht;

    PetscFunctionBegin;

    /* Used to keep track of the cells in the patch. */
    PetscHMapICreate(&ht);
    PetscHMapICreate(&cht);

    ierr = PCGetDM(pc, &dm); CHKERRQ(ierr);

    if (!dm) {
        SETERRQ(PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONGSTATE, "DM not yet set on patch PC\n");
    }

    ierr = PetscObjectTypeCompare((PetscObject)dm, DMPLEX, &flg); CHKERRQ(ierr);
    if (!flg) {
        SETERRQ(PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONGSTATE, "DM on patch PC must be DMPlex\n");
    }
    ierr = DMPlexGetChart(dm, &pStart, &pEnd); CHKERRQ(ierr);
    ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd); CHKERRQ(ierr);

    if (patch->user_patches) {
        /* compute patch->nuserIS, patch->userIS here */
        ierr = patch->userpatchconstructionop(pc, &patch->nuserIS, &patch->userIS, &patch->iterationSet, patch->userpatchconstructctx); CHKERRQ(ierr);
        vStart = 0;
        vEnd = patch->nuserIS;
    } else if (patch->codim < 0) { /* codim unset */
        if (patch->dim < 0) { /* dim unset */
            ierr = DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd); CHKERRQ(ierr);
        } else { /* dim set */
            ierr = DMPlexGetDepthStratum(dm, patch->dim, &vStart, &vEnd); CHKERRQ(ierr);
        }
    } else { /* codim set */
        ierr = DMPlexGetHeightStratum(dm, patch->codim, &vStart, &vEnd); CHKERRQ(ierr);
    }

    /* These labels mark the owned points.  We only create patches
     * around points that this process owns. */
    ierr = DMGetLabel(dm, "pyop2_ghost", &ghost); CHKERRQ(ierr);

    ierr = DMLabelCreateIndex(ghost, pStart, pEnd); CHKERRQ(ierr);

    ierr = PetscSectionCreate(PETSC_COMM_SELF, &patch->cellCounts); CHKERRQ(ierr);
    cellCounts = patch->cellCounts;
    ierr = PetscSectionSetChart(cellCounts, vStart, vEnd); CHKERRQ(ierr);

    /* Count cells in the patch surrounding each entity */
    for ( PetscInt v = vStart; v < vEnd; v++ ) {
        PetscHashIter hi;
        PetscInt chtSize;

        if (!patch->user_patches) {
            ierr = DMLabelHasPoint(ghost, v, &flg); CHKERRQ(ierr);
            /* Not an owned entity, don't make a cell patch. */
            if (flg) {
                continue;
            }
        }

        ierr = patch->patchconstructop((void*)patch, dm, v, ht); CHKERRQ(ierr);
        ierr = PCPatchCompleteCellPatch(dm, ht, cht);

        PetscHMapIGetSize(cht, &chtSize);
        if (chtSize == 0) {
            /* empty patch, continue */
            continue;
        }

        PetscHashIterBegin(cht, hi); /* safe because size(cht) > 0 from above */
        while (!PetscHashIterAtEnd(cht, hi)) {
            PetscInt entity;
            PetscHashIterGetKey(cht, hi, entity);
            if (cStart <= entity && entity < cEnd) {
                ierr = PetscSectionAddDof(cellCounts, v, 1); CHKERRQ(ierr);
            }
            PetscHashIterNext(cht, hi);
        }
    }
    ierr = DMLabelDestroyIndex(ghost); CHKERRQ(ierr);

    ierr = PetscSectionSetUp(cellCounts); CHKERRQ(ierr);
    ierr = PetscSectionGetStorageSize(cellCounts, &numCells); CHKERRQ(ierr);
    ierr = PetscMalloc1(numCells, &cellsArray); CHKERRQ(ierr);

    /* Now that we know how much space we need, run through again and
     * actually remember the cells. */
    for ( PetscInt v = vStart; v < vEnd; v++ ) {
        PetscInt ndof, off;
        PetscHashIter hi;

        ierr = PetscSectionGetDof(cellCounts, v, &ndof); CHKERRQ(ierr);
        ierr = PetscSectionGetOffset(cellCounts, v, &off); CHKERRQ(ierr);
        if ( ndof <= 0 ) {
            continue;
        }
        ierr = patch->patchconstructop((void*)patch, dm, v, ht); CHKERRQ(ierr);
        ierr = PCPatchCompleteCellPatch(dm, ht, cht);
        ndof = 0;
        PetscHashIterBegin(cht, hi);
        while (!PetscHashIterAtEnd(cht, hi)) {
            PetscInt entity;
            PetscHashIterGetKey(cht, hi, entity);
            if (cStart <= entity && entity < cEnd) {
                cellsArray[ndof + off] = entity;
                ndof++;
            }
            PetscHashIterNext(cht, hi);
        }
    }

    ierr = ISCreateGeneral(PETSC_COMM_SELF, numCells, cellsArray, PETSC_OWN_POINTER, &patch->cells); CHKERRQ(ierr);
    ierr = PetscSectionGetChart(patch->cellCounts, &pStart, &pEnd); CHKERRQ(ierr);
    patch->npatch = pEnd - pStart;
    PetscHMapIDestroy(&ht);
    PetscHMapIDestroy(&cht);
    PetscFunctionReturn(0);
}

/*
 * PCPatchCreateCellPatchDiscretisationInfo - Build the dof maps for cell patches
 *
 * Input Parameters:
 * + dm - The DMPlex object defining the mesh
 * . cellCounts - Section with counts of cells around each vertex
 * . cells - IS of the cell point indices of cells in each patch
 * . cellNumbering - Section mapping plex cell points to Firedrake cell indices.
 * . nodesPerCell - number of nodes per cell.
 * - cellNodeMap - map from cells to node indices (nodesPerCell * numCells)
 *
 * Output Parameters:
 * + dofs - IS of local dof numbers of each cell in the patch
 * . gtolCounts - Section with counts of dofs per cell patch
 * - gtol - IS mapping from global dofs to local dofs for each patch.
 */
static PetscErrorCode PCPatchCreateCellPatchDiscretisationInfo(PC pc)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch           = (PC_PATCH *)pc->data;
    PetscSection    cellCounts      = patch->cellCounts;
    PetscSection    gtolCounts;
    IS              cells           = patch->cells;
    PetscSection    cellNumbering   = patch->cellNumbering;
    PetscInt        numCells;
    PetscInt        numDofs;
    PetscInt        numGlobalDofs;
    PetscInt        totalDofsPerCell = patch->totalDofsPerCell;
    PetscInt        vStart, vEnd;
    const PetscInt *cellsArray;
    PetscInt       *newCellsArray   = NULL;
    PetscInt       *dofsArray       = NULL;
    PetscInt       *asmArray        = NULL;
    PetscInt       *globalDofsArray = NULL;
    PetscInt        globalIndex     = 0;
    PetscHMapI      ht;
    PetscHMapI      globalBcs;
    DM              dm         = NULL;
    const PetscInt *bcNodes    = NULL;
    PetscInt        numBcs;
    PetscHMapI      ownedpts, seenpts, owneddofs, seendofs, artificialbcs;
    PetscHashIter  hi;

    PetscFunctionBegin;

    ierr = PCGetDM(pc, &dm); CHKERRQ(ierr);

    /* dofcounts section is cellcounts section * dofPerCell */
    ierr = PetscSectionGetStorageSize(cellCounts, &numCells); CHKERRQ(ierr);
    numDofs = numCells * totalDofsPerCell;
    ierr = PetscMalloc1(numDofs, &dofsArray); CHKERRQ(ierr);
    ierr = PetscMalloc1(numDofs, &asmArray); CHKERRQ(ierr);
    ierr = PetscMalloc1(numCells, &newCellsArray); CHKERRQ(ierr);
    ierr = PetscSectionGetChart(cellCounts, &vStart, &vEnd); CHKERRQ(ierr);
    ierr = PetscSectionCreate(PETSC_COMM_SELF, &patch->gtolCounts); CHKERRQ(ierr);
    gtolCounts = patch->gtolCounts;
    ierr = PetscSectionSetChart(gtolCounts, vStart, vEnd); CHKERRQ(ierr);

    /* Outside the patch loop, get the dofs that are globally-enforced Dirichlet
       conditions */
    PetscHMapICreate(&globalBcs);
    ierr = ISGetIndices(patch->ghostBcNodes, &bcNodes); CHKERRQ(ierr);
    ierr = ISGetSize(patch->ghostBcNodes, &numBcs); CHKERRQ(ierr);
    for ( PetscInt i = 0; i < numBcs; i++ ) {
        PetscHMapISet(globalBcs, bcNodes[i], 0); /* these are already in concatenated numbering */
    }
    ierr = ISRestoreIndices(patch->ghostBcNodes, &bcNodes); CHKERRQ(ierr);

    /* HMap tables for artificial BC construction */
    PetscHMapICreate(&ownedpts);
    PetscHMapICreate(&seenpts);
    PetscHMapICreate(&owneddofs);
    PetscHMapICreate(&seendofs);
    PetscHMapICreate(&artificialbcs);

    ierr = ISGetIndices(cells, &cellsArray); CHKERRQ(ierr);
    PetscHMapICreate(&ht);
    for ( PetscInt v = vStart; v < vEnd; v++ ) {
        PetscInt dof, off;
        PetscInt localIndex = 0;
        PetscHMapIClear(ht);
        ierr = PetscSectionGetDof(cellCounts, v, &dof); CHKERRQ(ierr);
        ierr = PetscSectionGetOffset(cellCounts, v, &off); CHKERRQ(ierr);

        if ( dof <= 0 ) continue;

        /* Calculate the global numbers of the artificial BC dofs here first */
        ierr = patch->patchconstructop((void*)patch, dm, v, ownedpts); CHKERRQ(ierr);
        ierr = PCPatchCompleteCellPatch(dm, ownedpts, seenpts); CHKERRQ(ierr);
        ierr = PCPatchGetPointDofs(patch, ownedpts, owneddofs, v, patch->exclude_subspace); CHKERRQ(ierr);
        ierr = PCPatchGetPointDofs(patch, seenpts, seendofs, v, -1); CHKERRQ(ierr);
        ierr = PCPatchComputeSetDifference(owneddofs, seendofs, artificialbcs); CHKERRQ(ierr);
        if (patch->print_patches) {
            PetscHMapI globalbcdofs;
            PetscHMapICreate(&globalbcdofs);

            MPI_Comm comm = PetscObjectComm((PetscObject)pc);
            ierr = PetscSynchronizedPrintf(comm, "Patch %d: owned dofs:\n", v); CHKERRQ(ierr);
            PetscHashIterBegin(owneddofs, hi);
            while (!PetscHashIterAtEnd(owneddofs, hi)) {
                PetscInt globalDof;

                PetscHashIterGetKey(owneddofs, hi, globalDof);
                PetscHashIterNext(owneddofs, hi);
                ierr = PetscSynchronizedPrintf(comm, "%d ", globalDof); CHKERRQ(ierr);
            }
            ierr = PetscSynchronizedPrintf(comm, "\n"); CHKERRQ(ierr);
            ierr = PetscSynchronizedPrintf(comm, "Patch %d: seen dofs:\n", v); CHKERRQ(ierr);
            PetscHashIterBegin(seendofs, hi);
            while (!PetscHashIterAtEnd(seendofs, hi)) {
                PetscInt globalDof;
                PetscBool flg;

                PetscHashIterGetKey(seendofs, hi, globalDof);
                PetscHashIterNext(seendofs, hi);
                ierr = PetscSynchronizedPrintf(comm, "%d ", globalDof); CHKERRQ(ierr);

                PetscHMapIHas(globalBcs, globalDof, &flg);
                if (flg) {
                    PetscHMapISet(globalbcdofs, globalDof, 0);
                }
            }
            ierr = PetscSynchronizedPrintf(comm, "\n"); CHKERRQ(ierr);
            ierr = PetscSynchronizedPrintf(comm, "Patch %d: global BCs:\n", v); CHKERRQ(ierr);
            PetscHMapIGetSize(globalbcdofs, &numBcs);
            if (numBcs > 0) {
                PetscHashIterBegin(globalbcdofs, hi);
                while (!PetscHashIterAtEnd(globalbcdofs, hi)) {
                    PetscInt globalDof;
                    PetscHashIterGetKey(globalbcdofs, hi, globalDof);
                    PetscHashIterNext(globalbcdofs, hi);
                    ierr = PetscSynchronizedPrintf(comm, "%d ", globalDof); CHKERRQ(ierr);
                }
            }
            ierr = PetscSynchronizedPrintf(comm, "\n"); CHKERRQ(ierr);
            ierr = PetscSynchronizedPrintf(comm, "Patch %d: artificial BCs:\n", v); CHKERRQ(ierr);
            PetscHMapIGetSize(artificialbcs, &numBcs);
            if (numBcs > 0) {
                PetscHashIterBegin(artificialbcs, hi);
                while (!PetscHashIterAtEnd(artificialbcs, hi)) {
                    PetscInt globalDof;
                    PetscHashIterGetKey(artificialbcs, hi, globalDof);
                    PetscHashIterNext(artificialbcs, hi);
                    ierr = PetscSynchronizedPrintf(comm, "%d ", globalDof); CHKERRQ(ierr);
                }
            }
            ierr = PetscSynchronizedPrintf(comm, "\n\n"); CHKERRQ(ierr);
            PetscHMapIDestroy(&globalbcdofs);
        }
        for ( PetscInt k = 0; k < patch->nsubspaces; k++ ) {
            PetscInt nodesPerCell = patch->nodesPerCell[k];
            PetscInt subspaceOffset = patch->subspaceOffsets[k];
            const PetscInt *cellNodeMap = patch->cellNodeMap[k];
            PetscInt bs = patch->bs[k];

            for ( PetscInt i = off; i < off + dof; i++ ) {
                /* Walk over the cells in this patch. */
                const PetscInt c = cellsArray[i];
                PetscInt cell;
                ierr = PetscSectionGetDof(cellNumbering, c, &cell); CHKERRQ(ierr);
                if ( cell <= 0 ) {
                    SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE,
                            "Cell doesn't appear in cell numbering map");
                }
                ierr = PetscSectionGetOffset(cellNumbering, c, &cell); CHKERRQ(ierr);
                newCellsArray[i] = cell;
                for ( PetscInt j = 0; j < nodesPerCell; j++ ) {
                    /* For each global dof, map it into contiguous local storage. */
                    const PetscInt globalDof = cellNodeMap[cell*nodesPerCell + j]*bs + subspaceOffset;
                    /* finally, loop over block size */
                    for ( PetscInt l = 0; l < bs; l++ ) {
                        PetscInt localDof, isGlobalBcDof, isArtificialBcDof;

                        /* first, check if this is either a globally enforced or locally enforced BC dof */
                        PetscHMapIGet(globalBcs, globalDof + l, &isGlobalBcDof);
                        PetscHMapIGet(artificialbcs, globalDof + l, &isArtificialBcDof);

                        /* if it's either, don't ever give it a local dof number */
                        if (isGlobalBcDof >= 0 || isArtificialBcDof >= 0) {
                            dofsArray[globalIndex++] = -1; /* don't use this in assembly in this patch */
                        } else {
                            PetscHMapIGet(ht, globalDof + l, &localDof);
                            if (localDof == -1) {
                                localDof = localIndex++;
                                PetscHMapISet(ht, globalDof + l, localDof);
                            }
                            if ( globalIndex >= numDofs ) {
                                SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE,
                                        "Found more dofs than expected");
                            }
                            /* And store. */
                            dofsArray[globalIndex++] = localDof;
                        }
                    }
                }
            }
        }
        PetscHMapIGetSize(ht, &dof);
        /* How many local dofs in this patch? */
        ierr = PetscSectionSetDof(gtolCounts, v, dof); CHKERRQ(ierr);
    }
    if (globalIndex != numDofs) {
        SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE,
                 "Expected number of dofs (%d) doesn't match found number (%d)",
                 numDofs, globalIndex);
    }
    ierr = PetscSectionSetUp(gtolCounts); CHKERRQ(ierr);
    ierr = PetscSectionGetStorageSize(gtolCounts, &numGlobalDofs); CHKERRQ(ierr);
    ierr = PetscMalloc1(numGlobalDofs, &globalDofsArray); CHKERRQ(ierr);

    /* Now populate the global to local map.  This could be merged
    * into the above loop if we were willing to deal with reallocs. */
    PetscInt key = 0;
    PetscInt asmKey = 0;
    for ( PetscInt v = vStart; v < vEnd; v++ ) {
        PetscInt       dof, off;
        PetscHashIter hi;
        PetscHMapIClear(ht);
        ierr = PetscSectionGetDof(cellCounts, v, &dof); CHKERRQ(ierr);
        ierr = PetscSectionGetOffset(cellCounts, v, &off); CHKERRQ(ierr);

        if ( dof <= 0 ) continue;

        for ( PetscInt k = 0; k < patch->nsubspaces; k++ ) {
            PetscInt nodesPerCell = patch->nodesPerCell[k];
            PetscInt subspaceOffset = patch->subspaceOffsets[k];
            const PetscInt *cellNodeMap = patch->cellNodeMap[k];
            PetscInt bs = patch->bs[k];

            for ( PetscInt i = off; i < off + dof; i++ ) {
                /* Reconstruct mapping of global-to-local on this patch. */
                const PetscInt c = cellsArray[i];
                PetscInt cell;
                ierr = PetscSectionGetOffset(cellNumbering, c, &cell); CHKERRQ(ierr);
                for ( PetscInt j = 0; j < nodesPerCell; j++ ) {
                    for ( PetscInt l = 0; l < bs; l++ ) {
                        const PetscInt globalDof = cellNodeMap[cell*nodesPerCell + j]*bs + subspaceOffset + l;
                        const PetscInt localDof = dofsArray[key++];

                        if (localDof >= 0) PetscHMapISet(ht, globalDof, localDof);
                    }
                }
            }
            /* Shove it in the output data structure. */
            PetscInt goff;
            ierr = PetscSectionGetOffset(gtolCounts, v, &goff); CHKERRQ(ierr);
            PetscHashIterBegin(ht, hi);
            while (!PetscHashIterAtEnd(ht, hi)) {
                PetscInt globalDof, localDof;
                PetscHashIterGetKey(ht, hi, globalDof);
                PetscHashIterGetVal(ht, hi, localDof);
                if (globalDof >= 0) {
                    globalDofsArray[goff + localDof] = globalDof;
                }
                PetscHashIterNext(ht, hi);
            }
        }

        /* At this point, we have a hash table ht built that maps globalDof -> localDof.
           We need to create the dof table laid out cellwise first, then by subspace,
           as the assembler assembles cell-wise and we need to stuff the different
           contributions of the different function spaces to the right places. So we loop
           over cells, then over subspaces. */

        if (patch->nsubspaces > 1) { /* for nsubspaces = 1, data we need is already in dofsArray */
            for (PetscInt i = off; i < off + dof; i++ ) {
                const PetscInt c = cellsArray[i];
                PetscInt cell;
                ierr = PetscSectionGetOffset(cellNumbering, c, &cell); CHKERRQ(ierr);

                for ( PetscInt k = 0; k < patch->nsubspaces; k++ ) {
                    PetscInt nodesPerCell = patch->nodesPerCell[k];
                    PetscInt subspaceOffset = patch->subspaceOffsets[k];
                    const PetscInt *cellNodeMap = patch->cellNodeMap[k];
                    PetscInt bs = patch->bs[k];
                    for ( PetscInt j = 0; j < nodesPerCell; j++ ) {
                        for ( PetscInt l = 0; l < bs; l++ ) {
                            const PetscInt globalDof = cellNodeMap[cell*nodesPerCell + j]*bs + subspaceOffset + l;
                            PetscInt localDof;
                            PetscHMapIGet(ht, globalDof, &localDof);

                            /* If it's not in the hash table, i.e. is a BC dof,
                               then the PetscHMapIGet above gives -1, which matches
                               exactly the convention for PETSc's matrix assembly to
                               ignore the dof. So we don't need to do anything here */
                            asmArray[asmKey++] = localDof;
                        }
                    }
                }
            }
        }
    }

    if (1 == patch->nsubspaces) { /* replace with memcpy? */
        for (PetscInt i = 0; i < numDofs; i++) {
            asmArray[i] = dofsArray[i];
        }
    }

    PetscHMapIDestroy(&ht);
    ierr = ISRestoreIndices(cells, &cellsArray);
    ierr = PetscFree(dofsArray); CHKERRQ(ierr);

    /* Replace cell indices with firedrake-numbered ones. */
    ierr = ISGeneralSetIndices(cells, numCells, (const PetscInt *)newCellsArray, PETSC_OWN_POINTER); CHKERRQ(ierr);
    ierr = ISCreateGeneral(PETSC_COMM_SELF, numGlobalDofs, globalDofsArray, PETSC_OWN_POINTER, &patch->gtol); CHKERRQ(ierr);
    ierr = ISCreateGeneral(PETSC_COMM_SELF, numDofs, asmArray, PETSC_OWN_POINTER, &patch->dofs); CHKERRQ(ierr);

    PetscHMapIDestroy(&artificialbcs);
    PetscHMapIDestroy(&seendofs);
    PetscHMapIDestroy(&owneddofs);
    PetscHMapIDestroy(&seenpts);
    PetscHMapIDestroy(&ownedpts);
    PetscHMapIDestroy(&globalBcs);

    ierr = ISDestroy(&patch->ghostBcNodes); CHKERRQ(ierr); /* memory optimisation */
    PetscFunctionReturn(0);
}

static PetscErrorCode PCReset_PATCH(PC pc)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscInt        i;

    PetscFunctionBegin;
    ierr = PetscSFDestroy(&patch->defaultSF); CHKERRQ(ierr);
    ierr = PetscSectionDestroy(&patch->cellCounts); CHKERRQ(ierr);
    ierr = PetscSectionDestroy(&patch->cellNumbering); CHKERRQ(ierr);
    ierr = PetscSectionDestroy(&patch->gtolCounts); CHKERRQ(ierr);
    ierr = PetscSectionDestroy(&patch->bcCounts); CHKERRQ(ierr);
    ierr = ISDestroy(&patch->gtol); CHKERRQ(ierr);
    ierr = ISDestroy(&patch->cells); CHKERRQ(ierr);
    ierr = ISDestroy(&patch->dofs); CHKERRQ(ierr);
    ierr = ISDestroy(&patch->ghostBcNodes); CHKERRQ(ierr);
    ierr = ISDestroy(&patch->globalBcNodes); CHKERRQ(ierr);

    if (patch->dofSection) {
        for (i = 0; i < patch->nsubspaces; i++) {
            ierr = PetscSectionDestroy(&patch->dofSection[i]); CHKERRQ(ierr);
        }
    }
    ierr = PetscFree(patch->dofSection); CHKERRQ(ierr);
    ierr = PetscFree(patch->bs); CHKERRQ(ierr);
    ierr = PetscFree(patch->nodesPerCell); CHKERRQ(ierr);
    ierr = PetscFree(patch->cellNodeMap); CHKERRQ(ierr);
    ierr = PetscFree(patch->subspaceOffsets); CHKERRQ(ierr);

    if (patch->ksp) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = KSPReset(patch->ksp[i]); CHKERRQ(ierr);
        }
    }

    ierr = VecDestroy(&patch->localX); CHKERRQ(ierr);
    ierr = VecDestroy(&patch->localY); CHKERRQ(ierr);
    if (patch->patchX) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = VecDestroy(patch->patchX + i); CHKERRQ(ierr);
        }
        ierr = PetscFree(patch->patchX); CHKERRQ(ierr);
    }
    if (patch->patchY) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = VecDestroy(patch->patchY + i); CHKERRQ(ierr);
        }
        ierr = PetscFree(patch->patchY); CHKERRQ(ierr);
    }

    if (patch->partition_of_unity) {
        ierr = VecDestroy(&patch->dof_weights); CHKERRQ(ierr);
    }

    if (patch->patch_dof_weights) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = VecDestroy(patch->patch_dof_weights + i); CHKERRQ(ierr);
        }
        ierr = PetscFree(patch->patch_dof_weights); CHKERRQ(ierr);
    }
    if (patch->mat) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = MatDestroy(patch->mat + i); CHKERRQ(ierr);
        }
        ierr = PetscFree(patch->mat); CHKERRQ(ierr);
    }
    ierr = PetscFree(patch->sub_mat_type); CHKERRQ(ierr);

    patch->bs = 0;
    patch->cellNodeMap = NULL;

    if (patch->user_patches) {
        for ( PetscInt i = 0; i < patch->nuserIS; i++ ) {
            ierr = ISDestroy(&patch->userIS[i]); CHKERRQ(ierr);
        }
        PetscFree(patch->userIS);
        patch->nuserIS = 0;
    }
    if (patch->iterationSet) {
        ierr = ISDestroy(&patch->iterationSet); CHKERRQ(ierr);
    }

    PetscFunctionReturn(0);
}

static PetscErrorCode PCDestroy_PATCH(PC pc)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscInt        i;

    PetscFunctionBegin;

    ierr = PCReset_PATCH(pc); CHKERRQ(ierr);
    if (patch->ksp) {
        for ( i = 0; i < patch->npatch; i++ ) {
            ierr = KSPDestroy(&patch->ksp[i]); CHKERRQ(ierr);
        }
        ierr = PetscFree(patch->ksp); CHKERRQ(ierr);
    }
    ierr = PetscFree(pc->data); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCPatchZeroMatrix_Private(Mat mat, const PetscInt ncell,
                                                const PetscInt ndof,
                                                const PetscInt *dof)
{
    const PetscScalar *values = NULL;
    PetscInt rows;
    PetscErrorCode ierr;

    PetscFunctionBegin;

    ierr = PetscCalloc1(ndof*ndof, &values); CHKERRQ(ierr);
    for (PetscInt c = 0; c < ncell; c++) {
        const PetscInt *idx = &dof[ndof*c];
        ierr = MatSetValues(mat, ndof, idx, ndof, idx, values, INSERT_VALUES); CHKERRQ(ierr);
    }
    ierr = MatGetLocalSize(mat, &rows, NULL); CHKERRQ(ierr);
    for (PetscInt i = 0; i < rows; i++) {
        ierr = MatSetValues(mat, 1, &i, 1, &i, values, INSERT_VALUES); CHKERRQ(ierr);
    }
    ierr = MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = PetscFree(values); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCPatchCreateMatrix(PC pc, PetscInt which, Mat *mat)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscInt        csize, rsize;
    Vec             x, y;
    PetscBool       flg;
    const char     *prefix = NULL;

    PetscFunctionBegin;

    x = patch->patchX[which];
    y = patch->patchY[which];
    ierr = VecGetSize(x, &csize); CHKERRQ(ierr);
    ierr = VecGetSize(y, &rsize); CHKERRQ(ierr);
    ierr = MatCreate(PETSC_COMM_SELF, mat); CHKERRQ(ierr);
    ierr = PCGetOptionsPrefix(pc, &prefix); CHKERRQ(ierr);
    ierr = MatSetOptionsPrefix(*mat, prefix); CHKERRQ(ierr);
    ierr = MatAppendOptionsPrefix(*mat, "sub_"); CHKERRQ(ierr);
    if (patch->sub_mat_type) {
        ierr = MatSetType(*mat, patch->sub_mat_type); CHKERRQ(ierr);
    }
    ierr = MatSetSizes(*mat, rsize, csize, rsize, csize); CHKERRQ(ierr);
    ierr = PetscObjectTypeCompare((PetscObject)*mat, MATDENSE, &flg); CHKERRQ(ierr);
    if (!flg) {
        ierr = PetscObjectTypeCompare((PetscObject)*mat, MATSEQDENSE, &flg); CHKERRQ(ierr);
    }

    if (!flg) {
        PetscBT         bt;
        PetscInt       *dnnz       = NULL;
        const PetscInt *dofsArray = NULL;
        PetscInt        pStart, pEnd, ncell, offset;

        ierr = ISGetIndices(patch->dofs, &dofsArray); CHKERRQ(ierr);
        ierr = PetscSectionGetChart(patch->cellCounts, &pStart, &pEnd); CHKERRQ(ierr);

        which += pStart;
        if (which >= pEnd) {
            SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Asked for operator index is invalid\n"); CHKERRQ(ierr);
        }

        ierr = PetscSectionGetDof(patch->cellCounts, which, &ncell); CHKERRQ(ierr);
        ierr = PetscSectionGetOffset(patch->cellCounts, which, &offset); CHKERRQ(ierr);

        ierr = PetscMalloc1(rsize, &dnnz); CHKERRQ(ierr);
        for (PetscInt i = 0; i < rsize; i++) {
            dnnz[i] = 0;
        }
        ierr = PetscLogEventBegin(PC_Patch_Prealloc, pc, 0, 0, 0); CHKERRQ(ierr);

        /* XXX: This uses N^2 bits to store the sparsity pattern on a
         * patch.  This is probably OK if the patches are not too big,
         * but could use quite a bit of memory for planes in 3D.
         * Should we switch based on the value of rsize to a
         * hash-table (slower, but more memory efficient) approach? */
        ierr = PetscBTCreate(rsize*rsize, &bt); CHKERRQ(ierr);
        for (PetscInt c = 0; c < ncell; c++) {
            const PetscInt *idx = dofsArray + (offset + c)*patch->totalDofsPerCell;
            for (PetscInt i = 0; i < patch->totalDofsPerCell; i++) {
                const PetscInt row = idx[i];
                if (row < 0) continue;
                for (PetscInt j = 0; j < patch->totalDofsPerCell; j++) {
                    const PetscInt col = idx[j];
                    const PetscInt key = row*rsize + col;
                    if (col < 0) continue;

                    if (!PetscBTLookupSet(bt, key)) {
                        ++dnnz[row];
                    }
                }
            }
        }
        PetscBTDestroy(&bt);
        ierr = MatXAIJSetPreallocation(*mat, 1, dnnz, NULL, NULL, NULL); CHKERRQ(ierr);

        ierr = PetscFree(dnnz); CHKERRQ(ierr);
        ierr = PCPatchZeroMatrix_Private(*mat, ncell, patch->totalDofsPerCell,
                                         &dofsArray[offset*patch->totalDofsPerCell]); CHKERRQ(ierr);
        ierr = PetscLogEventEnd(PC_Patch_Prealloc, pc, 0, 0, 0); CHKERRQ(ierr);
        ierr = ISRestoreIndices(patch->dofs, &dofsArray); CHKERRQ(ierr);

    }

    ierr = MatSetUp(*mat); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

static PetscErrorCode PCPatchComputeOperator(PC pc, Mat mat, Mat multMat, PetscInt which)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    const PetscInt *dofsArray;
    const PetscInt *cellsArray;
    PetscInt        ncell, offset, pStart, pEnd;

    PetscFunctionBegin;

    ierr = PetscLogEventBegin(PC_Patch_ComputeOp, pc, 0, 0, 0); CHKERRQ(ierr);
    if (!patch->usercomputeop) {
        SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Must call PCPatchSetComputeOperator() to set user callback\n");
    }
    ierr = ISGetIndices(patch->dofs, &dofsArray); CHKERRQ(ierr);
    ierr = ISGetIndices(patch->cells, &cellsArray); CHKERRQ(ierr);
    ierr = PetscSectionGetChart(patch->cellCounts, &pStart, &pEnd); CHKERRQ(ierr);

    which += pStart;
    if (which >= pEnd) {
        SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Asked for operator index is invalid\n"); CHKERRQ(ierr);
    }

    ierr = PetscSectionGetDof(patch->cellCounts, which, &ncell); CHKERRQ(ierr);
    ierr = PetscSectionGetOffset(patch->cellCounts, which, &offset); CHKERRQ(ierr);
    if ( ncell <= 0 ) {
        ierr = PetscLogEventEnd(PC_Patch_ComputeOp, pc, 0, 0, 0); CHKERRQ(ierr);
        PetscFunctionReturn(0);
    }
    PetscStackPush("PCPatch user callback");
    ierr = patch->usercomputeop(pc, mat, ncell, cellsArray + offset, ncell*patch->totalDofsPerCell, dofsArray + offset*patch->totalDofsPerCell, patch->usercomputectx); CHKERRQ(ierr);
    PetscStackPop;
    ierr = ISRestoreIndices(patch->dofs, &dofsArray); CHKERRQ(ierr);
    ierr = ISRestoreIndices(patch->cells, &cellsArray); CHKERRQ(ierr);
    ierr = PetscLogEventEnd(PC_Patch_ComputeOp, pc, 0, 0, 0); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCPatch_ScatterLocal_Private(PC pc, PetscInt p,
                                                   Vec x, Vec y,
                                                   InsertMode mode,
                                                   ScatterMode scat)
{
    PetscErrorCode ierr;
    PC_PATCH          *patch   = (PC_PATCH *)pc->data;
    const PetscScalar *xArray = NULL;
    PetscScalar *yArray = NULL;
    const PetscInt *gtolArray = NULL;
    PetscInt offset, size;

    PetscFunctionBeginHot;
    ierr = PetscLogEventBegin(PC_Patch_Scatter, pc, 0, 0, 0); CHKERRQ(ierr);

    ierr = VecGetArrayRead(x, &xArray); CHKERRQ(ierr);
    ierr = VecGetArray(y, &yArray); CHKERRQ(ierr);

    ierr = PetscSectionGetDof(patch->gtolCounts, p, &size); CHKERRQ(ierr);
    ierr = PetscSectionGetOffset(patch->gtolCounts, p, &offset); CHKERRQ(ierr);
    ierr = ISGetIndices(patch->gtol, &gtolArray); CHKERRQ(ierr);
    if (mode == INSERT_VALUES && scat != SCATTER_FORWARD) {
        SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Can't insert if not scattering forward\n");
    }
    if (mode == ADD_VALUES && scat != SCATTER_REVERSE) {
        SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Can't add if not scattering reverse\n");
    }
    for ( PetscInt lidx = 0; lidx < size; lidx++ ) {
        const PetscInt gidx = gtolArray[lidx + offset];
        if (mode == INSERT_VALUES) {
            yArray[lidx] = xArray[gidx];
        } else {
            yArray[gidx] += xArray[lidx];
        }
    }
    ierr = VecRestoreArrayRead(x, &xArray); CHKERRQ(ierr);
    ierr = VecRestoreArray(y, &yArray); CHKERRQ(ierr);
    ierr = ISRestoreIndices(patch->gtol, &gtolArray); CHKERRQ(ierr);
    ierr = PetscLogEventEnd(PC_Patch_Scatter, pc, 0, 0, 0); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCSetUp_PATCH(PC pc)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    const char     *prefix;
    PetscInt       pStart;

    PetscFunctionBegin;

    if (!pc->setupcalled) {
        PetscInt     pStart, pEnd;
        PetscInt     localSize;
        ierr = PetscLogEventBegin(PC_Patch_CreatePatches, pc, 0, 0, 0); CHKERRQ(ierr);

        localSize = patch->subspaceOffsets[patch->nsubspaces];
        ierr = VecCreateSeq(PETSC_COMM_SELF, localSize, &patch->localX); CHKERRQ(ierr);
        ierr = VecSetUp(patch->localX); CHKERRQ(ierr);
        ierr = VecDuplicate(patch->localX, &patch->localY); CHKERRQ(ierr);
        ierr = PCPatchCreateCellPatches(pc); CHKERRQ(ierr);
        ierr = PCPatchCreateCellPatchDiscretisationInfo(pc); CHKERRQ(ierr);

        /* OK, now build the work vectors */
        ierr = PetscSectionGetChart(patch->gtolCounts, &pStart, &pEnd); CHKERRQ(ierr);
        ierr = PetscMalloc1(patch->npatch, &patch->patchX); CHKERRQ(ierr);
        ierr = PetscMalloc1(patch->npatch, &patch->patchY); CHKERRQ(ierr);

        for ( PetscInt i = pStart; i < pEnd; i++ ) {
            PetscInt dof;
            ierr = PetscSectionGetDof(patch->gtolCounts, i, &dof); CHKERRQ(ierr);
            ierr = VecCreateSeq(PETSC_COMM_SELF, dof, &patch->patchX[i - pStart]); CHKERRQ(ierr);
            ierr = VecSetUp(patch->patchX[i - pStart]); CHKERRQ(ierr);
            ierr = VecCreateSeq(PETSC_COMM_SELF, dof, &patch->patchY[i - pStart]); CHKERRQ(ierr);
            ierr = VecSetUp(patch->patchY[i - pStart]); CHKERRQ(ierr);
        }
        ierr = PetscMalloc1(patch->npatch, &patch->ksp); CHKERRQ(ierr);
        ierr = PCGetOptionsPrefix(pc, &prefix); CHKERRQ(ierr);
        for ( PetscInt i = 0; i < patch->npatch; i++ ) {
            ierr = KSPCreate(PETSC_COMM_SELF, patch->ksp + i); CHKERRQ(ierr);
            ierr = PetscObjectIncrementTabLevel((PetscObject)patch->ksp[i], (PetscObject)pc, 1); CHKERRQ(ierr);
            ierr = KSPSetOptionsPrefix(patch->ksp[i], prefix); CHKERRQ(ierr);
            ierr = KSPAppendOptionsPrefix(patch->ksp[i], "sub_"); CHKERRQ(ierr);
        }
        if (patch->save_operators) {
            ierr = PetscMalloc1(patch->npatch, &patch->mat); CHKERRQ(ierr);
            for ( PetscInt i = 0; i < patch->npatch; i++ ) {
                ierr = PCPatchCreateMatrix(pc, i, patch->mat + i); CHKERRQ(ierr);
            }
        }
        ierr = PetscLogEventEnd(PC_Patch_CreatePatches, pc, 0, 0, 0); CHKERRQ(ierr);
    }

    /* If desired, calculate weights for dof multiplicity */

    if (!pc->setupcalled && patch->partition_of_unity) {
        Mat P;
        Vec local;
        const PetscScalar *input = NULL;
        PetscScalar *output = NULL;
        ierr = PCGetOperators(pc, NULL, &P); CHKERRQ(ierr);
        ierr = MatCreateVecs(P, NULL, &patch->dof_weights);
        ierr = VecDuplicate(patch->localX, &local); CHKERRQ(ierr);
        ierr = PetscSectionGetChart(patch->gtolCounts, &pStart, NULL); CHKERRQ(ierr);
        for ( PetscInt i = 0; i < patch->npatch; i++ ) {
            PetscInt dof;
            ierr = PetscSectionGetDof(patch->gtolCounts, i + pStart, &dof); CHKERRQ(ierr);
            if ( dof <= 0 ) continue;
            ierr = VecSet(patch->patchX[i], 1.0); CHKERRQ(ierr);
            /* TODO: Do we need different scatters for X and Y? */

            ierr = PCPatch_ScatterLocal_Private(pc, i + pStart,
                                                patch->patchX[i], local,
                                                ADD_VALUES, SCATTER_REVERSE); CHKERRQ(ierr);
        }
        /* Local to Global */
        ierr = VecGetArrayRead(local, &input); CHKERRQ(ierr);
        ierr = VecGetArray(patch->dof_weights, &output); CHKERRQ(ierr);
        ierr = PetscSFReduceBegin(patch->defaultSF, MPIU_SCALAR, input, output, MPI_SUM); CHKERRQ(ierr);
        ierr = PetscSFReduceEnd(patch->defaultSF, MPIU_SCALAR, input, output, MPI_SUM); CHKERRQ(ierr);
        ierr = VecRestoreArray(patch->dof_weights, &output); CHKERRQ(ierr);
        ierr = VecRestoreArrayRead(local, &input); CHKERRQ(ierr);

        ierr = VecReciprocal(patch->dof_weights); CHKERRQ(ierr);
        ierr = VecDestroy(&local); CHKERRQ(ierr);
    }

    if (patch->save_operators) {
        for ( PetscInt i = 0; i < patch->npatch; i++ ) {
            ierr = MatZeroEntries(patch->mat[i]); CHKERRQ(ierr);
            ierr = PCPatchComputeOperator(pc, patch->mat[i], NULL, i); CHKERRQ(ierr);
            ierr = KSPSetOperators(patch->ksp[i], patch->mat[i], patch->mat[i]); CHKERRQ(ierr);
        }
    }
    if (!pc->setupcalled) {
        for ( PetscInt i = 0; i < patch->npatch; i++ ) {
            ierr = KSPSetFromOptions(patch->ksp[i]); CHKERRQ(ierr);
        }
    }
    PetscFunctionReturn(0);
}

static PetscErrorCode PCApply_PATCH(PC pc, Vec x, Vec y)
{
    PetscErrorCode     ierr;
    PC_PATCH          *patch   = (PC_PATCH *)pc->data;
    const PetscScalar *globalX = NULL;
    PetscScalar       *localX  = NULL;
    PetscScalar       *localY  = NULL;
    PetscScalar       *globalY = NULL;
    const PetscInt    *bcNodes = NULL;
    PetscInt           pStart, numBcs, size;
    PetscInt           nsweep = 1;
    PetscInt           start[2] = {0, patch->npatch-1};
    PetscInt           end[2] = {patch->npatch, -1};
    const PetscInt     inc[2] = {1, -1};
    const PetscInt    *iterationSet;
    PetscFunctionBegin;

    ierr = PetscLogEventBegin(PC_Patch_Apply, pc, 0, 0, 0); CHKERRQ(ierr);
    ierr = PetscOptionsPushGetViewerOff(PETSC_TRUE); CHKERRQ(ierr);
    /* Scatter from global space into overlapped local spaces */
    ierr = VecGetArrayRead(x, &globalX); CHKERRQ(ierr);
    ierr = VecGetArray(patch->localX, &localX); CHKERRQ(ierr);
    ierr = PetscSFBcastBegin(patch->defaultSF, MPIU_SCALAR, globalX, localX); CHKERRQ(ierr);
    ierr = PetscSFBcastEnd(patch->defaultSF, MPIU_SCALAR, globalX, localX); CHKERRQ(ierr);
    ierr = VecRestoreArrayRead(x, &globalX); CHKERRQ(ierr);
    ierr = VecRestoreArray(patch->localX, &localX); CHKERRQ(ierr);

    if (patch->user_patches) {
        ierr = ISGetLocalSize(patch->iterationSet, &end[0]); CHKERRQ(ierr);
        start[1] = end[0] - 1;
        ierr = ISGetIndices(patch->iterationSet, &iterationSet); CHKERRQ(ierr);
    }

    ierr = VecSet(patch->localY, 0.0); CHKERRQ(ierr);
    ierr = PetscSectionGetChart(patch->gtolCounts, &pStart, NULL); CHKERRQ(ierr);
    if (patch->symmetrise_sweep) {
        nsweep = 2;
    } else {
        nsweep = 1;
    }

    for (PetscInt sweep = 0; sweep < nsweep; sweep++) {
        for ( PetscInt j = start[sweep]; j*inc[sweep] < end[sweep]*inc[sweep]; j += inc[sweep] ) {
            PetscInt start, len, i;
            Mat multMat = NULL;

            if (patch->user_patches) {
                i = iterationSet[j];
            } else {
                i = j;
            }

            ierr = PetscSectionGetDof(patch->gtolCounts, i + pStart, &len); CHKERRQ(ierr);
            ierr = PetscSectionGetOffset(patch->gtolCounts, i + pStart, &start); CHKERRQ(ierr);
            if ( len <= 0 ) {
                /* TODO: Squash out these guys in the setup as well. */
                continue;
            }
            ierr = PCPatch_ScatterLocal_Private(pc, i + pStart,
                                                patch->localX, patch->patchX[i],
                                                INSERT_VALUES,
                                                SCATTER_FORWARD); CHKERRQ(ierr);
            if (!patch->save_operators) {
                Mat mat;
                ierr = PCPatchCreateMatrix(pc, i, &mat); CHKERRQ(ierr);
                /* Populate operator here. */
                ierr = PCPatchComputeOperator(pc, mat, multMat, i); CHKERRQ(ierr);
                ierr = KSPSetOperators(patch->ksp[i], mat, mat);
                /* Drop reference so the KSPSetOperators below will blow it away. */
                ierr = MatDestroy(&mat); CHKERRQ(ierr);
            }

            ierr = PetscLogEventBegin(PC_Patch_Solve, pc, 0, 0, 0); CHKERRQ(ierr);
            ierr = KSPSolve(patch->ksp[i], patch->patchX[i], patch->patchY[i]); CHKERRQ(ierr);
            ierr = PetscLogEventEnd(PC_Patch_Solve, pc, 0, 0, 0); CHKERRQ(ierr);
            if (!patch->save_operators) {
                PC pc;
                ierr = KSPSetOperators(patch->ksp[i], NULL, NULL); CHKERRQ(ierr);
                ierr = KSPGetPC(patch->ksp[i], &pc); CHKERRQ(ierr);
                /* Destroy PC context too, otherwise the factored matrix hangs around. */
                ierr = PCReset(pc); CHKERRQ(ierr);
            }

            ierr = PCPatch_ScatterLocal_Private(pc, i + pStart,
                                                patch->patchY[i], patch->localY,
                                                ADD_VALUES, SCATTER_REVERSE); CHKERRQ(ierr);
        }
    }

    if (patch->user_patches) {
        ierr = ISRestoreIndices(patch->iterationSet, &iterationSet); CHKERRQ(ierr);
    }

    /* Now patch->localY contains the solution of the patch solves, so
     * we need to combine them all. */
    ierr = VecSet(y, 0.0); CHKERRQ(ierr);
    ierr = VecGetArray(y, &globalY); CHKERRQ(ierr);
    ierr = VecGetArrayRead(patch->localY, (const PetscScalar **)&localY); CHKERRQ(ierr);
    ierr = PetscSFReduceBegin(patch->defaultSF, MPIU_SCALAR, localY, globalY, MPI_SUM); CHKERRQ(ierr);
    ierr = PetscSFReduceEnd(patch->defaultSF, MPIU_SCALAR, localY, globalY, MPI_SUM); CHKERRQ(ierr);
    ierr = VecRestoreArrayRead(patch->localY, (const PetscScalar **)&localY); CHKERRQ(ierr);

    /* Now we need to send the global BC values through */
    ierr = VecGetArrayRead(x, &globalX); CHKERRQ(ierr);
    ierr = ISGetSize(patch->globalBcNodes, &numBcs); CHKERRQ(ierr);
    ierr = ISGetIndices(patch->globalBcNodes, &bcNodes); CHKERRQ(ierr);
    ierr = VecGetLocalSize(x, &size); CHKERRQ(ierr);
    for ( PetscInt i = 0; i < numBcs; i++ ) {
        const PetscInt idx = bcNodes[i];
        if (idx < size) {
            globalY[idx] = globalX[idx];
        }
    }

    ierr = ISRestoreIndices(patch->globalBcNodes, &bcNodes); CHKERRQ(ierr);
    ierr = VecRestoreArrayRead(x, &globalX); CHKERRQ(ierr);
    ierr = VecRestoreArray(y, &globalY); CHKERRQ(ierr);
    if (patch->partition_of_unity) {
        /* Now apply partition of unity */
        ierr = VecPointwiseMult(y, y, patch->dof_weights); CHKERRQ(ierr);
    }

    ierr = PetscOptionsPopGetViewerOff(); CHKERRQ(ierr);
    ierr = PetscLogEventEnd(PC_Patch_Apply, pc, 0, 0, 0); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCSetUpOnBlocks_PATCH(PC pc)
{
  PC_PATCH           *patch = (PC_PATCH*)pc->data;
  PetscErrorCode      ierr;
  PetscInt            i;
  KSPConvergedReason  reason;

  PetscFunctionBegin;
  PetscFunctionReturn(0);
  for (i=0; i<patch->npatch; i++) {
    ierr = KSPSetUp(patch->ksp[i]); CHKERRQ(ierr);
    ierr = KSPGetConvergedReason(patch->ksp[i], &reason); CHKERRQ(ierr);
    if (reason == KSP_DIVERGED_PCSETUP_FAILED) {
      pc->failedreason = PC_SUBPC_ERROR;
    }
  }
  PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchConstruct_Star(void *vpatch, DM dm, PetscInt entity, PetscHMapI ht)
{
    PetscErrorCode ierr;
    PetscInt       starSize;
    PetscInt      *star = NULL;

    PetscFunctionBegin;
    PetscHMapIClear(ht);

    /* To start with, add the entity we care about */
    PetscHMapISet(ht, entity, 0);

    /* Loop over all the points that this entity connects to */
    ierr = DMPlexGetTransitiveClosure(dm, entity, PETSC_FALSE, &starSize, &star); CHKERRQ(ierr);
    for ( PetscInt si = 0; si < starSize; si++ ) {
        PetscInt pt = star[2*si];
        PetscHMapISet(ht, pt, 0);
    }
    if (star) {
        ierr = DMPlexRestoreTransitiveClosure(dm, 0, PETSC_FALSE, NULL, &star); CHKERRQ(ierr);
    }
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCPatchConstruct_Vanka(void *vpatch, DM dm, PetscInt entity, PetscHMapI ht)
{
    PetscErrorCode ierr;
    PC_PATCH      *patch = (PC_PATCH*) vpatch;
    PetscInt       starSize, closureSize;
    PetscInt      *star = NULL, *closure = NULL;
    PetscInt       iStart, iEnd;
    PetscInt       cStart, cEnd;
    PetscBool      shouldIgnore;

    PetscFunctionBegin;
    PetscHMapIClear(ht);

    /* To start with, add the entity we care about */
    PetscHMapISet(ht, entity, 0);

    /* Should we ignore any topological entities of a certain dimension? */
    if (patch->vankadim >= 0) {
        shouldIgnore = PETSC_TRUE;
        ierr = DMPlexGetDepthStratum(dm, patch->vankadim, &iStart, &iEnd); CHKERRQ(ierr);
    } else {
        shouldIgnore = PETSC_FALSE;
    }
    ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd); CHKERRQ(ierr);

    /* Loop over all the cells that this entity connects to */
    ierr = DMPlexGetTransitiveClosure(dm, entity, PETSC_FALSE, &starSize, &star); CHKERRQ(ierr);
    for ( PetscInt si = 0; si < starSize; si++ ) {
        PetscInt cell = star[2*si];
        if ( cell < cStart || cell >= cEnd) continue;
        /* now loop over all entities in the closure of that cell */
        ierr = DMPlexGetTransitiveClosure(dm, cell, PETSC_TRUE, &closureSize, &closure); CHKERRQ(ierr);
        for ( PetscInt ci = 0; ci < closureSize; ci++ ) {
            PetscInt newentity = closure[2*ci];
            if (shouldIgnore && iStart <= newentity && newentity < iEnd) {
                /* We've been told to ignore entities of this type.*/
                continue;
            }
            PetscHMapISet(ht, newentity, 0);
        }
    }
    if (closure) {
        ierr = DMPlexRestoreTransitiveClosure(dm, 0, PETSC_TRUE, NULL, &closure); CHKERRQ(ierr);
    }
    if (star) {
        ierr = DMPlexRestoreTransitiveClosure(dm, 0, PETSC_FALSE, NULL, &star); CHKERRQ(ierr);
    }
    PetscFunctionReturn(0);
}

/* The user's already set the patches in patch->userIS. Build the hash tables */
PETSC_EXTERN PetscErrorCode PCPatchConstruct_User(void *vpatch, DM dm, PetscInt entity, PetscHMapI ht)
{
    PetscErrorCode  ierr;
    PC_PATCH       *patch = (PC_PATCH*) vpatch;
    IS              patchis = patch->userIS[entity];
    PetscInt        size;
    PetscInt        pStart, pEnd;
    const PetscInt *patchdata;

    PetscFunctionBegin;
    PetscHMapIClear(ht);

    ierr = DMPlexGetChart(dm, &pStart, &pEnd);

    ierr = ISGetLocalSize(patchis, &size); CHKERRQ(ierr);
    ierr = ISGetIndices(patchis, &patchdata); CHKERRQ(ierr);
    for ( PetscInt i = 0; i < size; i++ ) {
        PetscInt ownedentity = patchdata[i];
        if (ownedentity < pStart || ownedentity >= pEnd) {
            SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_USER,"Entities need to be between the bounds of DMPlexGetChart()");
        }
        PetscHMapISet(ht, patchdata[i], 0);
    }
    ierr = ISRestoreIndices(patchis, &patchdata); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

const char *const PCPatchConstructTypes[]   = {"star","vanka","user","python",0};

static PetscErrorCode PCSetFromOptions_PATCH(PetscOptionItems *PetscOptionsObject, PC pc)
{
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscErrorCode  ierr;
    PetscBool       flg, dimflg, codimflg;
    char            sub_mat_type[256];
    PCPatchConstructType patchConstructionType = PC_PATCH_STAR;

    PetscFunctionBegin;
    ierr = PetscOptionsHead(PetscOptionsObject, "Vertex-patch Additive Schwarz options"); CHKERRQ(ierr);

    ierr = PetscOptionsBool("-pc_patch_save_operators", "Store all patch operators for lifetime of PC?",
                            "PCPatchSetSaveOperators", patch->save_operators, &patch->save_operators, &flg); CHKERRQ(ierr);

    ierr = PetscOptionsBool("-pc_patch_partition_of_unity", "Weight contributions by dof multiplicity?",
                            "PCPatchSetPartitionOfUnity", patch->partition_of_unity, &patch->partition_of_unity, &flg); CHKERRQ(ierr);

    ierr = PetscOptionsBool("-pc_patch_multiplicative", "Gauss-Seidel instead of Jacobi?",
                            "PCPatchSetMultiplicative", patch->multiplicative, &patch->multiplicative, &flg); CHKERRQ(ierr);
    if (patch->multiplicative) {
      SETERRQ(PetscObjectComm((PetscObject)pc),PETSC_ERR_USER,"We've removed multiplicative to do BC condensation (for now)");
    }

    ierr = PetscOptionsInt("-pc_patch_construction_dim", "What dimension of entity to construct patches by? (0 = vertices)", "PCSetFromOptions_PATCH", patch->dim, &patch->dim, &dimflg);
    ierr = PetscOptionsInt("-pc_patch_construction_codim", "What co-dimension of entity to construct patches by? (0 = cells)", "PCSetFromOptions_PATCH", patch->codim, &patch->codim, &codimflg);
    if (dimflg && codimflg) {
        SETERRQ(PetscObjectComm((PetscObject)pc),PETSC_ERR_USER,"Can only set one of dimension or co-dimension");
    }
    /* XXX: This should be PetscOptionsEnum */
    ierr = PetscOptionsEList("-pc_patch_construction_type", "How should the patches be constructed?", "PCSetFromOptions_PATCH", PCPatchConstructTypes, 4, PCPatchConstructTypes[patchConstructionType], (PetscInt *)&patchConstructionType, &flg);
    if (flg) {
        switch (patchConstructionType) {
        case PC_PATCH_STAR:
            patch->patchconstructop = PCPatchConstruct_Star;
            break;
        case PC_PATCH_VANKA:
            patch->patchconstructop = PCPatchConstruct_Vanka;
            ierr = PetscOptionsInt("-pc_patch_vanka_dim", "Topological dimension of entities for Vanka to ignore", "PCSetFromOptions_PATCH", patch->vankadim, &patch->vankadim, &flg);
            ierr = PetscOptionsInt("-pc_patch_vanka_space", "What subspace is the constraint space for Vanka?", "PCSetFromOptions_PATCH", patch->exclude_subspace, &patch->exclude_subspace, &flg);
            if (flg) {
                SETERRQ(PetscObjectComm((PetscObject)pc),PETSC_ERR_USER, "-pc_patch_vanka_space has been renamed to -pc_patch_exclude_subspace");
            }
            break;
        case PC_PATCH_USER:
        case PC_PATCH_PYTHON:
            patch->user_patches = PETSC_TRUE;
            patch->patchconstructop = PCPatchConstruct_User;
            break;
        default:
            SETERRQ(PetscObjectComm((PetscObject)pc),PETSC_ERR_USER,"Unknown patch construction type");
            break;
        }
    }

    ierr = PetscOptionsFList("-pc_patch_sub_mat_type", "Matrix type for patch solves", "PCPatchSetSubMatType",MatList, NULL, sub_mat_type, 256, &flg); CHKERRQ(ierr);
    if (flg) {
        ierr = PCPatchSetSubMatType(pc, sub_mat_type); CHKERRQ(ierr);
    }

    ierr = PetscOptionsBool("-pc_patch_print_patches", "Print out information during patch construction?",
                            "PCSetFromOptions_PATCH", patch->print_patches, &patch->print_patches, &flg); CHKERRQ(ierr);

    ierr = PetscOptionsBool("-pc_patch_symmetrise_sweep", "Go start->end, end->start?",
                            "PCSetFromOptions_PATCH", patch->symmetrise_sweep, &patch->symmetrise_sweep, &flg); CHKERRQ(ierr);

    ierr = PetscOptionsInt("-pc_patch_exclude_subspace", "What subspace (if any) to exclude in construction?", "PCSetFromOptions_PATCH", patch->exclude_subspace, &patch->exclude_subspace, &flg);

    ierr = PetscOptionsTail(); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

static PetscErrorCode PCView_PATCH(PC pc, PetscViewer viewer)
{
    PC_PATCH       *patch = (PC_PATCH *)pc->data;
    PetscErrorCode  ierr;
    PetscMPIInt     rank;
    PetscBool       isascii;
    PetscViewer     sviewer;
    PetscFunctionBegin;
    ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&isascii);CHKERRQ(ierr);

    ierr = MPI_Comm_rank(PetscObjectComm((PetscObject)pc),&rank);CHKERRQ(ierr);
    if (!isascii) {
        PetscFunctionReturn(0);
    }
    ierr = PetscViewerASCIIPushTab(viewer); CHKERRQ(ierr);
    ierr = PetscViewerASCIIPrintf(viewer, "Subspace Correction preconditioner with %d patches\n", patch->npatch); CHKERRQ(ierr);
    if (patch->multiplicative) {
        ierr = PetscViewerASCIIPrintf(viewer, "Schwarz type: multiplicative\n"); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPrintf(viewer, "Schwarz type: additive\n"); CHKERRQ(ierr);
    }
    if (patch->partition_of_unity) {
        ierr = PetscViewerASCIIPrintf(viewer, "Weighting by partition of unity\n"); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPrintf(viewer, "Not weighting by partition of unity\n"); CHKERRQ(ierr);
    }
    if (patch->symmetrise_sweep) {
        ierr = PetscViewerASCIIPrintf(viewer, "Symmetrising sweep (start->end, then end->start)\n"); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPrintf(viewer, "Not symmetrising sweep\n"); CHKERRQ(ierr);
    }
    if (!patch->save_operators) {
        ierr = PetscViewerASCIIPrintf(viewer, "Not saving patch operators (rebuilt every PCApply)\n"); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPrintf(viewer, "Saving patch operators (rebuilt every PCSetUp)\n"); CHKERRQ(ierr);
    }
    if (patch->patchconstructop == PCPatchConstruct_Star) {
        ierr = PetscViewerASCIIPrintf(viewer, "Patch construction operator: star\n"); CHKERRQ(ierr);
    } else if (patch->patchconstructop == PCPatchConstruct_Vanka) {
        ierr = PetscViewerASCIIPrintf(viewer, "Patch construction operator: Vanka\n"); CHKERRQ(ierr);
    } else if (patch->patchconstructop == PCPatchConstruct_User) {
        ierr = PetscViewerASCIIPrintf(viewer, "Patch construction operator: user-specified\n"); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPrintf(viewer, "Patch construction operator: unknown\n"); CHKERRQ(ierr);
    }
    ierr = PetscViewerASCIIPrintf(viewer, "KSP on patches (all same):\n"); CHKERRQ(ierr);

    if (patch->ksp) {
        ierr = PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer); CHKERRQ(ierr);
        if (!rank) {
            ierr = PetscViewerASCIIPushTab(sviewer); CHKERRQ(ierr);
            ierr = KSPView(patch->ksp[0], sviewer); CHKERRQ(ierr);
            ierr = PetscViewerASCIIPopTab(sviewer); CHKERRQ(ierr);
        }
        ierr = PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer); CHKERRQ(ierr);
    } else {
        ierr = PetscViewerASCIIPushTab(viewer); CHKERRQ(ierr);
        ierr = PetscViewerASCIIPrintf(viewer, "KSP not yet set.\n"); CHKERRQ(ierr);
        ierr = PetscViewerASCIIPopTab(viewer); CHKERRQ(ierr);
    }

    ierr = PetscViewerASCIIPopTab(viewer); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PETSC_EXTERN PetscErrorCode PCCreate_PATCH(PC pc)
{
    PetscErrorCode ierr;
    PC_PATCH       *patch;

    PetscFunctionBegin;

    ierr = PetscNewLog(pc, &patch); CHKERRQ(ierr);

    /* Set some defaults */
    patch->sub_mat_type       = NULL;
    patch->save_operators     = PETSC_TRUE;
    patch->partition_of_unity = PETSC_FALSE;
    patch->multiplicative     = PETSC_FALSE;
    patch->codim              = -1;
    patch->dim                = -1;
    patch->exclude_subspace   = -1;
    patch->vankadim           = -1;
    patch->patchconstructop   = PCPatchConstruct_Star;
    patch->print_patches      = PETSC_FALSE;
    patch->symmetrise_sweep   = PETSC_FALSE;
    patch->nuserIS            = 0;
    patch->userIS             = NULL;
    patch->iterationSet       = NULL;
    patch->user_patches       = PETSC_FALSE;

    pc->data                 = (void *)patch;
    pc->ops->apply           = PCApply_PATCH;
    pc->ops->applytranspose  = 0; /* PCApplyTranspose_PATCH; */
    pc->ops->setup           = PCSetUp_PATCH;
    pc->ops->reset           = PCReset_PATCH;
    pc->ops->destroy         = PCDestroy_PATCH;
    pc->ops->setfromoptions  = PCSetFromOptions_PATCH;
    pc->ops->setuponblocks   = PCSetUpOnBlocks_PATCH;
    pc->ops->view            = PCView_PATCH;
    pc->ops->applyrichardson = 0;

    PetscFunctionReturn(0);
}

ext: c

max_stars_repo_name: wence-/ssc

max_issues_repo_name: wence-/ssc

max_forks_repo_name: wence-/ssc


text:
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <unistd.h>

#ifdef GSL_FOUND
#include <gsl/gsl_integration.h>
#endif

#include "core_allvars.h"
#include "core_init.h"
#include "core_mymalloc.h"
#include "core_cool_func.h"

/* These functions do not need to be exposed externally */
double integrand_time_to_present(const double a, void *param);
void set_units(struct params *run_params);
void read_snap_list(const int ThisTask, struct params *run_params);
double time_to_present(const double z, struct params *run_params);

#ifdef HDF5
#include "io/io_save_hdf5.h"
#endif

void init(const int ThisTask, struct params *run_params)
{
    run_params->Age = mymalloc(ABSOLUTEMAXSNAPS*sizeof(run_params->Age[0]));

    set_units(run_params);

    read_snap_list(ThisTask, run_params);

    //Hack to fix deltaT for snapshot 0
    //This way, galsnapnum = -1 will not segfault.
    run_params->Age[0] = time_to_present(1000.0, run_params);//lookback time from z=1000
    run_params->Age++;

    for(int i = 0; i < run_params->Snaplistlen; i++) {
        run_params->ZZ[i] = 1 / run_params->AA[i] - 1;
        run_params->Age[i] = time_to_present(run_params->ZZ[i], run_params);
    }

    run_params->a0 = 1.0 / (1.0 + run_params->Reionization_z0);
    run_params->ar = 1.0 / (1.0 + run_params->Reionization_zr);

    read_cooling_functions();
    if(ThisTask == 0) {
        printf("cooling functions read\n\n");
    }

#if 0
#ifdef HDF5
    if(HDF5Output) {
        calc_hdf5_props();
    }
#endif
#endif

}

void set_units(struct params *run_params)
{

    run_params->UnitTime_in_s = run_params->UnitLength_in_cm / run_params->UnitVelocity_in_cm_per_s;
    run_params->UnitTime_in_Megayears = run_params->UnitTime_in_s / SEC_PER_MEGAYEAR;
    run_params->G = GRAVITY / pow(run_params->UnitLength_in_cm, 3) * run_params->UnitMass_in_g * pow(run_params->UnitTime_in_s, 2);
    run_params->UnitDensity_in_cgs = run_params->UnitMass_in_g / pow(run_params->UnitLength_in_cm, 3);
    run_params->UnitPressure_in_cgs = run_params->UnitMass_in_g / run_params->UnitLength_in_cm / pow(run_params->UnitTime_in_s, 2);
    run_params->UnitCoolingRate_in_cgs = run_params->UnitPressure_in_cgs / run_params->UnitTime_in_s;
    run_params->UnitEnergy_in_cgs = run_params->UnitMass_in_g * pow(run_params->UnitLength_in_cm, 2) / pow(run_params->UnitTime_in_s, 2);

    run_params->EnergySNcode = run_params->EnergySN / run_params->UnitEnergy_in_cgs * run_params->Hubble_h;
    run_params->EtaSNcode = run_params->EtaSN * (run_params->UnitMass_in_g / SOLAR_MASS) / run_params->Hubble_h;

    // convert some physical input parameters to internal units
    run_params->Hubble = HUBBLE * run_params->UnitTime_in_s;

    // compute a few quantitites
    run_params->RhoCrit = 3.0 * run_params->Hubble * run_params->Hubble / (8 * M_PI * run_params->G);
}

void read_snap_list(const int ThisTask, struct params *run_params)
{
    char fname[MAX_STRING_LEN+1];

    snprintf(fname, MAX_STRING_LEN, "%s", run_params->FileWithSnapList);
    FILE *fd = fopen(fname, "r");
    if(fd == NULL) {
        printf("can't read output list in file '%s'\n", fname);
        ABORT(0);
    }

    run_params->Snaplistlen = 0;
    do {
        if(fscanf(fd, " %lg ", &(run_params->AA[run_params->Snaplistlen])) == 1) {
            run_params->Snaplistlen++;
        } else {
            break;
        }
    } while(run_params->Snaplistlen < run_params->MAXSNAPS);
    fclose(fd);

    if(ThisTask == 0) {
        printf("found %d defined times in snaplist\n", run_params->Snaplistlen);
    }
}

double time_to_present(const double z, struct params *run_params)
{
    const double end_limit = 1.0;
    const double start_limit = 1.0/(1 + z);
    double result=0.0;
#ifdef GSL_FOUND
#define WORKSIZE 1000
    gsl_function F;
    gsl_integration_workspace *workspace;
    double abserr;

    workspace = gsl_integration_workspace_alloc(WORKSIZE);
    F.function = &integrand_time_to_present;
    F.params = run_params;

    gsl_integration_qag(&F, start_limit, end_limit, 1.0 / run_params->Hubble,
                        1.0e-9, WORKSIZE, GSL_INTEG_GAUSS21, workspace, &result, &abserr);

    gsl_integration_workspace_free(workspace);

#undef WORKSIZE
#else
    /* Do not have GSL - let's integrate numerically ourselves */
    const double step  = 1e-7;
    const int64_t nsteps = (end_limit - start_limit)/step;
    result = 0.0;
    const double y0 = integrand_time_to_present(start_limit + 0*step, run_params);
    const double yn = integrand_time_to_present(start_limit + nsteps*step, run_params);
    for(int64_t i=1; i<nsteps; i++) {
        result  += integrand_time_to_present(start_limit + i*step, run_params);
    }

    result = (step*0.5)*(y0 + yn + 2.0*result);
#endif

    /* convert into Myrs/h (I think -> MS 23/6/2018) */
    const double time = 1.0 / run_params->Hubble * result;

    // return time to present as a function of redshift
    return time;
}

double integrand_time_to_present(const double a, void *param)
{
    const struct params *run_params = (struct params *) param;
    return 1.0 / sqrt(run_params->Omega / a + (1.0 - run_params->Omega - run_params->OmegaLambda) + run_params->OmegaLambda * a * a);
}


text:
#ifndef COMUTILS_H
#define    COMUTILS_H

#include "HypCommands.h" // RECORD_LENGTH

#include <gsl/gsl-lite.hpp> // gsl::span

#include <array>
#include <functional>
#include <memory>
#include <stdint.h>
#include <string>
#include <vector>

#include <QObject>
#include <QString>
QT_USE_NAMESPACE

class SerialPort;

#define val const auto

//------------------------------------------------------------------------------
class DeviceLink : public QObject
{
    Q_OBJECT

    DeviceLink()                              = delete;
    DeviceLink            (const DeviceLink&) = delete;
    DeviceLink& operator =(const DeviceLink&) = delete;

public:
    //enum class DeviceType {
    //    UNKNOWN = 0,
    //    LDU, RDU, MDU,
    //    NUM_TYPES // the last
    //};

    typedef std::function<void(const QString&)>            MsgToDisplayCbk;
    typedef std::function<void(const gsl::span<uint8_t>&)> PacketRecvdCbk;
    typedef std::function<void()>                          SeriesEndCbk;
    typedef std::function<void(bool)>                      FinishedCbk;

    DeviceLink(MsgToDisplayCbk msgCbk, PacketRecvdCbk packetCbk,
               SeriesEndCbk seriesEndCbk, FinishedCbk finishCbk);

    void DownloadRecorded(); // returns packets via PacketRecvdCbk
    void ClearRecordings();
    void UploadToDevice(const std::vector<uint8_t>& data);

    Hyperion::DeviceType GetDeviceType() const { return m_deviceType;  }
    int GetFirmwareVersionX100() const { return m_firmwareVersionX100; }

private:
    void DownloadThreadFn();
    void ClearThreadFn();
    void UploadThreadFn();

    std::unique_ptr<SerialPort> EstablishConnection(size_t& port_idx, uint8_t& handshakeReply);

    void SendMessage  (const QString& msg) const { if (m_messageCbk)     m_messageCbk(msg); }
    void ReceivePacket(const gsl::span<uint8_t>& p)const { if (m_packetRecvdCbk) m_packetRecvdCbk(p); }
    void SendFinished (bool ok)            const { if (m_finishCbk)      m_finishCbk(ok); }

    typedef std::vector<uint8_t> dataVec_t;

    template<size_t N> static bool ChecksumOk(const std::array<uint8_t, N>& data) {
        size_t sum = 0;
        for (size_t i=0;  i+1 < N;  sum += data[i++]);
        return ((sum & 0xFF) == data[N - 1]);
    }

//data:
    MsgToDisplayCbk m_messageCbk;
    PacketRecvdCbk  m_packetRecvdCbk;
    SeriesEndCbk    m_seriesEndCbk;
    FinishedCbk     m_finishCbk;

    std::vector<QString> m_portNames;

    Hyperion::DeviceType m_deviceType = Hyperion::UNKNOWN_DEVICE;
    int m_firmwareVersionX100 = 0; // unknown

//statics const:
    static const std::array<uint8_t,2> HANDSHAKE;            // "DZ"
    static const std::array<uint8_t,2> START_DOWNLOAD;       // "DK"
    static const std::array<uint8_t,1> CONTINUE_DOWNLOAD;    // 0x01
    static const std::array<uint8_t,2> CLEAR_VALUES;         // "DE" - not sure what it does...
    static const std::array<uint8_t,2> CLEAR_MEMORY;         // "DB"
    static const std::array<uint8_t,2> GET_FW_VERSION; // "DV"

    static const std::array<uint8_t,5> HANDSHAKE_REPLY;
};
//------------------------------------------------------------------------------

#endif // COMUTILS_H

ext: h


text:
#pragma once

#include <memory>
#include <string>
#include <unordered_map>

#include <DirectXMath.h>
#include <gsl/span>

struct Ray3d;

namespace gfx {

	struct Light3d;
	struct MdfRenderOverrides;
	enum class MaterialPlaceholderSlot;

	struct AnimatedModelParams;
	using AnimatedModelPtr = std::shared_ptr<class AnimatedModel>;
	using MdfRenderMaterialPtr = std::shared_ptr<class MdfRenderMaterial>;

	enum class WeaponAnim : int {
		None = 0,
		RightAttack,
		RightAttack2,
		RightAttack3,
		LeftAttack,
		LeftAttack2,
		LeftAttack3,
		Walk,
		Run,
		Idle,
		FrontHit,
		FrontHit2,
		FrontHit3,
		LeftHit,
		LeftHit2,
		LeftHit3,
		RightHit,
		RightHit2,
		RightHit3,
		BackHit,
		BackHit2,
		BackHit3,
		RightCriticalSwing,
		LeftCriticalSwing,
		Fidget,
		Fidget2,
		Fidget3,
		Sneak,
		Panic,
		RightCombatStart,
		LeftCombatStart,
		CombatIdle,
		CombatFidget,
		Special1,
		Special2,
		Special3,
		FrontDodge,
		RightDodge,
		LeftDodge,
		BackDodge,
		RightThrow,
		LeftThrow,
		LeftSnatch,
		RightSnatch,
		LeftTurn,
		RightTurn
	};

	enum class WeaponAnimType : int {
		Unarmed = 0,
		Dagger,
		Sword,
		Mace,
		Hammer,
		Axe,
		Club,
		Battleaxe,
		Greatsword,
		Greataxe,
		Greathammer,
		Spear,
		Staff,
		Polearm,
		Bow,
		Crossbow,
		Sling,
		Shield,
		Flail,
		Chain,
		TwoHandedFlail,
		Shuriken,
		Monk
	};

	enum class BardInstrumentType : int {
		Flute = 0,
		Drum,
		Mandolin,
		Trumpet,
		Harp,
		Lute,
		Pipers,
		Recorder
	};

	enum class NormalAnimType : int {
		Falldown = 0,
		ProneIdle,
		ProneFidget,
		Getup,
		Magichands,
		Picklock,
		PicklockConcentrated,
		Examine,
		Throw,
		Death,
		Death2,
		Death3,
		DeadIdle,
		DeadFidget,
		DeathProneIdle,
		DeathProneFidget,
		AbjurationCasting,
		AbjurationConjuring,
		ConjurationCasting,
		ConjurationConjuring,
		DivinationCasting,
		DivinationConjuring,
		EnchantmentCasting,
		EnchantmentConjuring,
		EvocationCasting,
		EvocationConjuring,
		IllusionCasting,
		IllusionConjuring,
		NecromancyCasting,
		NecromancyConjuring,
		TransmutationCasting,
		TransmutationConjuring,
		Conceal,
		ConcealIdle,
		Unconceal,
		ItemIdle,
		ItemFidget,
		Open,
		Close,
		SkillAnimalEmpathy,
		SkillDisableDevice,
		SkillHeal,
		SkillHealConcentrated,
		SkillHide,
		SkillHideIdle,
		SkillHideFidget,
		SkillUnhide,
		SkillPickpocket,
		SkillSearch,
		SkillSpot,
		FeatTrack,
		Trip,
		Bullrush,
		Flurry,
		Kistrike,
		Tumble,
		Special1,
		Special2,
		Special3,
		Special4,
		Throw2,
		WandAbjurationCasting,
		WandAbjurationConjuring,
		WandConjurationCasting,
		WandConjurationConjuring,
		WandDivinationCasting,
		WandDivinationConjuring,
		WandEnchantmentCasting,
		WandEnchantmentConjuring,
		WandEvocationCasting,
		WandEvocationConjuring,
		WandIllusionCasting,
		WandIllusionConjuring,
		WandNecromancyCasting,
		WandNecromancyConjuring,
		WandTransmutationCasting,
		WandTransmutationConjuring,
		SkillBarbarianRage,
		OpenIdle
	};

	/*
	Represents an encoded animation id.
	*/
	class EncodedAnimId {
	public:
		explicit EncodedAnimId(int id) : mId(id) {
		}

		explicit EncodedAnimId(WeaponAnim anim,
		                       WeaponAnimType leftHand = WeaponAnimType::Unarmed,
		                       WeaponAnimType rightHand = WeaponAnimType::Unarmed) : mId(sWeaponAnimFlag) {
			auto animId = (int)anim;
			auto leftHandId = (int)leftHand;
			auto rightHandId = (int)rightHand;

			mId |= animId & 0xFFFFF;
			mId |= leftHandId << 20;
			mId |= rightHandId << 25;
		}

		explicit EncodedAnimId(BardInstrumentType instrumentType) : mId(sBardInstrumentAnimFlag) {
			mId |= (int)instrumentType;
		}

		explicit EncodedAnimId(NormalAnimType animType) : mId((int) animType) {
		}

		operator int() const {
			return mId;
		}

		bool IsConjuireAnimation() const;

		bool IsSpecialAnim() const {
			return (mId & (sWeaponAnimFlag | sBardInstrumentAnimFlag)) != 0;
		}

		// Not for weapon/bard anims
		NormalAnimType GetNormalAnimType() const {
			return (NormalAnimType)mId;
		}

		bool IsWeaponAnim() const {
			return (mId & sWeaponAnimFlag) != 0;
		}

		// Only valid for weapon animations
		WeaponAnimType GetWeaponLeftHand() const {
			return (WeaponAnimType)((mId >> 20) & 0x1F);
		}

		// Only valid for weapon animations
		WeaponAnimType GetWeaponRightHand() const {
			return (WeaponAnimType)((mId >> 25) & 0x1F);
		}

		// Only valid for weapon animations
		WeaponAnim GetWeaponAnim() const {
			return (WeaponAnim)(mId & 0xFFFFF);
		}

		bool IsBardInstrumentAnim() const {
			return (mId & sBardInstrumentAnimFlag) != 0;
		}

		// Only valid for bard instrument anim
		BardInstrumentType GetBardInstrumentType() const {
			return (BardInstrumentType)(mId & 7);
		}

		std::string GetName() const;

		bool ToFallback();

	private:
		// Indicates that an animation id uses the encoded format
		static constexpr int sWeaponAnimFlag = 1 << 30;
		static constexpr int sBardInstrumentAnimFlag = 1 << 31;

		int mId;
	};

	/*
		Represents the events that can trigger when the animation
		of an animated model is advanced.
	*/
	class AnimatedModelEvents {
	public:
		AnimatedModelEvents(bool isEnd, bool isAction)
			: mIsEnd(isEnd),
			  mIsAction(isAction) {
		}

		bool IsEnd() const {
			return mIsEnd;
		}

		bool IsAction() const {
			return mIsAction;
		}

	private:
		const bool mIsEnd : 1;
		const bool mIsAction : 1;
	};

	class Submesh {
	public:
		virtual ~Submesh() {
		}

		virtual int GetVertexCount() = 0;
		virtual int GetPrimitiveCount() = 0;
		virtual gsl::span<DirectX::XMFLOAT4> GetPositions() = 0;
		virtual gsl::span<DirectX::XMFLOAT4> GetNormals() = 0;
		virtual gsl::span<DirectX::XMFLOAT2> GetUV() = 0;
		virtual gsl::span<uint16_t> GetIndices() = 0;
	};

	class IRenderState {
	public:
		virtual ~IRenderState() = default;
	};

	class AnimatedModel {
	public:
		virtual ~AnimatedModel() {
		}

		virtual uint32_t GetHandle() const = 0;

		virtual bool AddAddMesh(const std::string& filename) = 0;

		virtual bool ClearAddMeshes() = 0;

		virtual AnimatedModelEvents Advance(float deltaTime,
		                                    float deltaDistance,
		                                    float deltaRotation,
		                                    const AnimatedModelParams& params) = 0;

		virtual EncodedAnimId GetAnimId() const = 0;

		virtual int GetBoneCount() const = 0;

		virtual std::string GetBoneName(int boneId) = 0;

		virtual int GetBoneParentId(int boneId) = 0;

		virtual bool GetBoneWorldMatrixByName(
			const AnimatedModelParams& params,
			const std::string& boneName,
			DirectX::XMFLOAT4X4* worldMatrixOut) = 0;

		virtual bool GetBoneWorldMatrixByNameForChild(const AnimatedModelPtr& child,
		                                              const AnimatedModelParams& params,
		                                              const std::string& boneName,
				DirectX::XMFLOAT4X4* worldMatrixOut) = 0;

		virtual float GetDistPerSec() const = 0;

		virtual float GetRotationPerSec() const = 0;

		virtual bool HasAnim(EncodedAnimId animId) const = 0;

		virtual void SetTime(const AnimatedModelParams& params, float timeInSecs) = 0;

		virtual bool HasBone(const std::string& boneName) const = 0;

		virtual void AddReplacementMaterial(gfx::MaterialPlaceholderSlot slot,
			const gfx::MdfRenderMaterialPtr &material) = 0;

		virtual void SetAnimId(EncodedAnimId animId) = 0;

		// This seems to reset cloth simulation state
		virtual void SetClothFlag() = 0;

		virtual std::vector<int> GetSubmeshes() = 0;

		virtual std::unique_ptr<Submesh> GetSubmesh(const AnimatedModelParams& params, int submeshIdx) = 0;

		virtual std::unique_ptr<Submesh> GetSubmeshForParticles(const AnimatedModelParams& params, int submeshIdx) = 0;

		bool HitTestRay(const AnimatedModelParams& params, const Ray3d &ray, float &hitDistance);

		/**
		 * Find the closest distance that the given point is away from the surface of this mesh.
		 */
		float GetDistanceToMesh(const AnimatedModelParams &params, DirectX::XMFLOAT3 pos);

		/**
			This calculates the effective height in world coordinate units of the model in its current
			state. Scale is the model scale in percent.
		*/
		virtual float GetHeight(int scale = 100) = 0;

		/**
			This calculates the visible radius of the model in its current state.
			The radius is the maximum distance of any vertex on the x,z plane from the models origin.
			If the model has no vertices, 0 is returned.
			Scale is model scale in percent.
		*/
		virtual float GetRadius(int scale = 100) = 0;

		/**
		 * Sets a custom render state pointer that will be freed when this model is freed.
		 */
		virtual void SetRenderState(std::unique_ptr<IRenderState> renderState) = 0;

		/**
		 * Returns the currently assigned render state or null.
		 */
		virtual IRenderState *GetRenderState() const = 0;

		template<typename T>
		T &GetOrCreateRenderState() {
			auto current = GetRenderState();
			if (!current) {
				auto newState = std::make_unique<T>();
				current = newState.get();
				SetRenderState(std::move(newState));
			}
			return (T&) *current;
		}

	};

	struct AnimatedModelParams { // see: objects.GetAnimParams(handle)
		uint32_t x = 0;
		uint32_t y = 0;
		float offsetX = 0;
		float offsetY = 0;
		float offsetZ = 0;
		float rotation = 0;
		float scale = 1;
		float rotationRoll = 0;
		float rotationPitch = 0;
		float rotationYaw = 0;
		AnimatedModelPtr parentAnim;
		std::string attachedBoneName;
		bool rotation3d = false; // Enables use of rotationRoll/rotationPitch/rotationYaw
	};

	class AnimatedModelFactory {
	public:
		virtual ~AnimatedModelFactory() {
		}

		virtual AnimatedModelPtr FromIds(
			int meshId,
			int skeletonId,
			EncodedAnimId idleAnimId,
			const AnimatedModelParams& params,
			bool borrow = false) = 0;

		virtual AnimatedModelPtr FromFilenames(
			const std::string& meshFilename,
			const std::string& skeletonFilename,
			EncodedAnimId idleAnimId,
			const AnimatedModelParams& params) = 0;

		virtual std::unique_ptr<gfx::AnimatedModel> BorrowByHandle(uint32_t handle) = 0;

		virtual void FreeHandle(uint32_t handle) = 0;

		virtual void FreeAll() = 0;

	};

	class AnimatedModelRenderer {
	public:
		virtual ~AnimatedModelRenderer() {}

		virtual void Render(AnimatedModel *model,
			const AnimatedModelParams& params,
			gsl::span<Light3d> lights,
			const MdfRenderOverrides *materialOverrides = nullptr) = 0;
	};

}

max_stars_repo_name: edoipi/TemplePlus

max_issues_repo_name: edoipi/TemplePlus

max_forks_repo_name: edoipi/TemplePlus


text:
#include "../include/matrix.h"

#ifndef MATREX_NO_BLAS

#include <cblas.h>

void
matrix_dot(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

void
matrix_dot_and_add(
  const float alpha, const Matrix first, const Matrix second, const Matrix third, Matrix result
) {
  const uint64_t data_size = MX_ROWS(first) * MX_COLS(second) + 2;

  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );

  for(uint64_t index = 2; index < data_size; index += 1) {
    result[index] += third[index];
  }
}

void
matrix_dot_and_apply(
  const float alpha, const Matrix first, const Matrix second, const char *function_name, Matrix result
) {
  const math_func_ptr_t func = math_func_from_name(function_name);

  const uint64_t data_size = MX_ROWS(first) * MX_COLS(second) + 2;

  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );

  for(uint64_t index = 2; index < data_size; index += 1) {
    result[index] = func(result[index]);
  }
}

void
matrix_dot_nt(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_ROWS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasTrans,
    MX_ROWS(first),
    MX_ROWS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

void
matrix_dot_tn(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_COLS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasTrans,
    CblasNoTrans,
    MX_COLS(first),
    MX_COLS(second),
    MX_ROWS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

#else

void
matrix_dot(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_and_add(
  const float alpha, const Matrix first, const Matrix second, const Matrix third, Matrix result
) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = third[elem_offset];
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_and_apply(
  const float alpha, const Matrix first, const Matrix second, const char *function_name, Matrix result
) {
  const math_func_ptr_t func = math_func_from_name(function_name);

  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] = func(alpha * result[elem_offset]);
    }
}

void
matrix_dot_nt(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_ROWS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + c*MX_COLS(second) + k];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_tn(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_COLS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_ROWS(first); k++)
        result[elem_offset] += first[2 + r + k*rows]*second[2 + c + k*cols];
      result[elem_offset] *= alpha;
    }
}

#endif

max_stars_repo_name: scripbox/matrex

max_issues_repo_name: scripbox/matrex

max_forks_repo_name: scripbox/matrex