WIP: draft of mutation metadata schema

fwdpy11/tskit_tools/metadata_schema.py

@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
 #
+import copy
 import typing
 
 import fwdpy11
@@ -60,6 +61,39 @@
     }
 )
 
+MutationMetadata = tskit.metadata.MetadataSchema(
+    {
+        "codec": "struct",
+        "type": "object",
+        "name": "Mutation metadata",
+        "properties": {
+            "s": {"type": "number", "binaryFormat": "d"},
+            "h": {"type": "number", "binaryFormat": "d"},
+            "origin": {"type": "number", "binaryFormat": "I"},
+            "neutral": {"type": "number", "binaryFormat": "?"},
+            "label": {"type": "number", "binaryFormat": "H"},
+            "key": {"type": "number", "binaryFormat": "Q"},
+        },
+        "additionalProperties": False,
+    }
+)
+
+_MutationMetaWithVectorsDict = copy.deepcopy(MutationMetadata.schema)
+
+_MutationMetaWithVectorsDict["name"] = "Mutation metadata with vectors"
+_MutationMetaWithVectorsDict["properties"]["esizes"] = {
+    "type": "array",
+    "items": {"type": "number", "binaryFormat": "d"},
+}
+_MutationMetaWithVectorsDict["properties"]["heffects"] = {
+    "type": "array",
+    "items": {"type": "number", "binaryFormat": "d"},
+}
+
+MutationMetadataWithVectors = tskit.metadata.MetadataSchema(
+    _MutationMetaWithVectorsDict
+)
+
 
 def generate_individual_metadata(
     metadata: fwdpy11._fwdpy11.DiploidMetadata,
@@ -76,3 +110,12 @@ def generate_individual_metadata(
         "label": metadata.label,
     }
     return d
+
+
+def determine_mutation_metadata_schema(
+    mutations: fwdpy11._fwdpy11.MutationVector,
+) -> tskit.metadata.MetadataSchema:
+    if len(mutations) == 0 or len(mutations[0].esizes) == 0:
+        return MutationMetadata
+
+    return MutationMetadataWithVectors