supporting branching tensors (without physical axes) in a TTNO

qc-tum · Oct 14, 2024 · 300bc2d · 300bc2d
1 parent 2b9a31d
commit 300bc2d
Show file tree

Hide file tree

Showing 12 changed files with 504 additions and 339 deletions.
diff --git a/src/operator/local_op.c b/src/operator/local_op.c
@@ -50,54 +50,121 @@ void construct_local_operator(const struct local_op_ref* opics, const int nopics
 {
 	assert(nopics > 0);
 
-	// first summand
-	copy_dense_tensor(&opmap[opics[0].oid], op);
-	switch (op->dtype)
+	if (opics[0].oid == OID_NOP)
 	{
-		case CT_SINGLE_REAL:
+		// construct dummy identity operator
+		const long dim[2] = { 1, 1 };
+		allocate_dense_tensor(opmap[0].dtype, 2, dim, op);
+		dense_tensor_set_identity(op);
+		// scale by weight factor
+		switch (op->dtype)
 		{
-			const float* cmap = coeffmap;
-			scale_dense_tensor(&cmap[opics[0].cid], op);
-			// add the other summands
-			for (int i = 1; i < nopics; i++) {
-				dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+			case CT_SINGLE_REAL:
+			{
+				const float* cmap = coeffmap;
+				float weight = 0;
+				for (int i = 0; i < nopics; i++) {
+					// operators must be consistent
+					assert(opics[i].oid == OID_NOP);
+					weight += cmap[opics[i].cid];
+				}
+				scale_dense_tensor(&weight, op);
+				break;
 			}
-			break;
-		}
-		case CT_DOUBLE_REAL:
-		{
-			const double* cmap = coeffmap;
-			scale_dense_tensor(&cmap[opics[0].cid], op);
-			// add the other summands
-			for (int i = 1; i < nopics; i++) {
-				dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+			case CT_DOUBLE_REAL:
+			{
+				const double* cmap = coeffmap;
+				double weight = 0;
+				for (int i = 0; i < nopics; i++) {
+					// operators must be consistent
+					assert(opics[i].oid == OID_NOP);
+					weight += cmap[opics[i].cid];
+				}
+				scale_dense_tensor(&weight, op);
+				break;
 			}
-			break;
-		}
-		case CT_SINGLE_COMPLEX:
-		{
-			const scomplex* cmap = coeffmap;
-			scale_dense_tensor(&cmap[opics[0].cid], op);
-			// add the other summands
-			for (int i = 1; i < nopics; i++) {
-				dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+			case CT_SINGLE_COMPLEX:
+			{
+				const scomplex* cmap = coeffmap;
+				scomplex weight = 0;
+				for (int i = 0; i < nopics; i++) {
+					// operators must be consistent
+					assert(opics[i].oid == OID_NOP);
+					weight += cmap[opics[i].cid];
+				}
+				scale_dense_tensor(&weight, op);
+				break;
 			}
-			break;
-		}
-		case CT_DOUBLE_COMPLEX:
-		{
-			const dcomplex* cmap = coeffmap;
-			scale_dense_tensor(&cmap[opics[0].cid], op);
-			// add the other summands
-			for (int i = 1; i < nopics; i++) {
-				dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+			case CT_DOUBLE_COMPLEX:
+			{
+				const dcomplex* cmap = coeffmap;
+				dcomplex weight = 0;
+				for (int i = 0; i < nopics; i++) {
+					// operators must be consistent
+					assert(opics[i].oid == OID_NOP);
+					weight += cmap[opics[i].cid];
+				}
+				scale_dense_tensor(&weight, op);
+				break;
+			}
+			default:
+			{
+				// unknown data type
+				assert(false);
 			}
-			break;
 		}
-		default:
+	}
+	else
+	{
+		// first summand
+		copy_dense_tensor(&opmap[opics[0].oid], op);
+		switch (op->dtype)
 		{
-			// unknown data type
-			assert(false);
+			case CT_SINGLE_REAL:
+			{
+				const float* cmap = coeffmap;
+				scale_dense_tensor(&cmap[opics[0].cid], op);
+				// add the other summands
+				for (int i = 1; i < nopics; i++) {
+					dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+				}
+				break;
+			}
+			case CT_DOUBLE_REAL:
+			{
+				const double* cmap = coeffmap;
+				scale_dense_tensor(&cmap[opics[0].cid], op);
+				// add the other summands
+				for (int i = 1; i < nopics; i++) {
+					dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+				}
+				break;
+			}
+			case CT_SINGLE_COMPLEX:
+			{
+				const scomplex* cmap = coeffmap;
+				scale_dense_tensor(&cmap[opics[0].cid], op);
+				// add the other summands
+				for (int i = 1; i < nopics; i++) {
+					dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+				}
+				break;
+			}
+			case CT_DOUBLE_COMPLEX:
+			{
+				const dcomplex* cmap = coeffmap;
+				scale_dense_tensor(&cmap[opics[0].cid], op);
+				// add the other summands
+				for (int i = 1; i < nopics; i++) {
+					dense_tensor_scalar_multiply_add(&cmap[opics[i].cid], &opmap[opics[i].oid], op);
+				}
+				break;
+			}
+			default:
+			{
+				// unknown data type
+				assert(false);
+			}
 		}
 	}
 }
diff --git a/src/operator/local_op.h b/src/operator/local_op.h
@@ -12,7 +12,8 @@
 ///
 enum
 {
-	OID_IDENTITY = 0,  //!< index corresponding to identity operation
+	OID_NOP      = -1,  //!< index corresponding to no operation
+	OID_IDENTITY =  0,  //!< index corresponding to identity operation
 };