1: #include <petscis.h>
  2: #include <petsc/private/isimpl.h>
  3: #include <petsc/private/viewerimpl.h>
  4: #include <petsclayouthdf5.h>

  6: PetscErrorCode ISView_Binary(IS is, PetscViewer viewer)
  7: {
  8:   PetscBool       skipHeader;
  9:   PetscLayout     map;
 10:   PetscInt        tr[2], n, s, N;
 11:   const PetscInt *iarray;

 13:   PetscFunctionBegin;
 14:   PetscCall(PetscViewerSetUp(viewer));
 15:   PetscCall(PetscViewerBinaryGetSkipHeader(viewer, &skipHeader));

 17:   PetscCall(ISGetLayout(is, &map));
 18:   PetscCall(PetscLayoutGetLocalSize(map, &n));
 19:   PetscCall(PetscLayoutGetRange(map, &s, NULL));
 20:   PetscCall(PetscLayoutGetSize(map, &N));

 22:   /* write IS header */
 23:   tr[0] = IS_FILE_CLASSID;
 24:   tr[1] = N;
 25:   if (!skipHeader) PetscCall(PetscViewerBinaryWrite(viewer, tr, 2, PETSC_INT));

 27:   /* write IS indices */
 28:   PetscCall(ISGetIndices(is, &iarray));
 29:   PetscCall(PetscViewerBinaryWriteAll(viewer, iarray, n, s, N, PETSC_INT));
 30:   PetscCall(ISRestoreIndices(is, &iarray));
 31:   PetscFunctionReturn(PETSC_SUCCESS);
 32: }

 34: #if PetscDefined(HAVE_HDF5)
 35: /*
 36:      This should handle properly the cases where PetscInt is 32 or 64 and hsize_t is 32 or 64. That is properly casting with
 37:    checks back and forth between the two types of variables.
 38: */
 39: static PetscErrorCode ISLoad_HDF5(IS is, PetscViewer viewer)
 40: {
 41:   PetscInt   *ind;
 42:   const char *isname;

 44:   PetscFunctionBegin;
 45:   PetscCheck(((PetscObject)is)->name, PetscObjectComm((PetscObject)is), PETSC_ERR_SUP, "IS name must be given using PetscObjectSetName() before ISLoad() since HDF5 can store multiple objects in a single file");
 46:   PetscCall(PetscObjectGetName((PetscObject)is, &isname));
 47:   PetscCall(PetscViewerHDF5Load(viewer, isname, is->map, PetscDefined(USE_64BIT_INDICES) ? H5T_NATIVE_LLONG : H5T_NATIVE_INT, (void **)&ind));
 48:   PetscCall(ISGeneralSetIndices(is, is->map->n, ind, PETSC_OWN_POINTER));
 49:   PetscCall(PetscInfo(is, "Read IS object with name %s of size %" PetscInt_FMT ":%" PetscInt_FMT "\n", isname, is->map->n, is->map->N));
 50:   PetscFunctionReturn(PETSC_SUCCESS);
 51: }
 52: #endif

 54: static PetscErrorCode ISLoad_Binary(IS is, PetscViewer viewer)
 55: {
 56:   PetscBool   isgeneral, skipHeader;
 57:   PetscInt    tr[2], rows, N, n, s, *idx;
 58:   PetscLayout map;

 60:   PetscFunctionBegin;
 61:   PetscCall(PetscObjectTypeCompare((PetscObject)is, ISGENERAL, &isgeneral));
 62:   PetscCheck(isgeneral, PetscObjectComm((PetscObject)is), PETSC_ERR_ARG_INCOMP, "IS must be of type ISGENERAL to load into it");
 63:   PetscCall(PetscViewerSetUp(viewer));
 64:   PetscCall(PetscViewerBinaryGetSkipHeader(viewer, &skipHeader));

 66:   PetscCall(ISGetLayout(is, &map));
 67:   PetscCall(PetscLayoutGetSize(map, &N));

 69:   /* read IS header */
 70:   if (!skipHeader) {
 71:     PetscCall(PetscViewerBinaryRead(viewer, tr, 2, NULL, PETSC_INT));
 72:     PetscCheck(tr[0] == IS_FILE_CLASSID, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Not an IS next in file");
 73:     PetscCheck(tr[1] >= 0, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "IS size (%" PetscInt_FMT ") in file is negative", tr[1]);
 74:     PetscCheck(N < 0 || N == tr[1], PETSC_COMM_SELF, PETSC_ERR_FILE_UNEXPECTED, "IS in file different size (%" PetscInt_FMT ") than input IS (%" PetscInt_FMT ")", tr[1], N);
 75:     rows = tr[1];
 76:   } else {
 77:     PetscCheck(N >= 0, PETSC_COMM_SELF, PETSC_ERR_USER, "IS binary file header was skipped, thus the user must specify the global size of input IS");
 78:     rows = N;
 79:   }

 81:   /* set IS size if not already set */
 82:   if (N < 0) PetscCall(PetscLayoutSetSize(map, rows));
 83:   PetscCall(PetscLayoutSetUp(map));

 85:   /* get IS sizes and check global size */
 86:   PetscCall(PetscLayoutGetSize(map, &N));
 87:   PetscCall(PetscLayoutGetLocalSize(map, &n));
 88:   PetscCall(PetscLayoutGetRange(map, &s, NULL));
 89:   PetscCheck(N == rows, PETSC_COMM_SELF, PETSC_ERR_FILE_UNEXPECTED, "IS in file different size (%" PetscInt_FMT ") than input IS (%" PetscInt_FMT ")", rows, N);

 91:   /* read IS indices */
 92:   PetscCall(PetscMalloc1(n, &idx));
 93:   PetscCall(PetscViewerBinaryReadAll(viewer, idx, n, s, N, PETSC_INT));
 94:   PetscCall(ISGeneralSetIndices(is, n, idx, PETSC_OWN_POINTER));
 95:   PetscFunctionReturn(PETSC_SUCCESS);
 96: }

 98: PetscErrorCode ISLoad_Default(IS is, PetscViewer viewer)
 99: {
100:   PetscBool isbinary, ishdf5;

102:   PetscFunctionBegin;
103:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
104:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
105:   if (isbinary) {
106:     PetscCall(ISLoad_Binary(is, viewer));
107:   } else if (ishdf5) {
108: #if PetscDefined(HAVE_HDF5)
109:     PetscCall(ISLoad_HDF5(is, viewer));
110: #endif
111:   }
112:   PetscFunctionReturn(PETSC_SUCCESS);
113: }