Support reshaping with correct dimensions from NDArray

ophyd/v2/_p4p.py

@@ -3,9 +3,10 @@
 from asyncio import CancelledError
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, Optional, Sequence, Type, Union
+from typing import Any, Dict, List, Optional, Sequence, Type, Union
 
 from bluesky.protocols import Descriptor, Dtype, Reading
+from p4p import Value
 from p4p.client.asyncio import Context, Subscription
 
 from .core import (
@@ -55,12 +56,54 @@ def descriptor(self, source: str, value) -> Descriptor:
         dtype = specifier_to_dtype[value.type().aspy("value")]
         return dict(source=source, dtype=dtype, shape=[])
 
+    def metadata_fields(self) -> List[str]:
+        """
+        PVA request string for metadata.
+        """
+        return ["alarm", "timeStamp"]
+
+    def value_fields(self) -> List[str]:
+        """
+        PVA request string for value only.
+        """
+        return ["value"]
+
 
 class PvaArrayConverter(PvaConverter):
     def descriptor(self, source: str, value) -> Descriptor:
         return dict(source=source, dtype="array", shape=[len(value["value"])])
 
 
+class PvaNDArrayConverter(PvaConverter):
+    def metadata_fields(self) -> List[str]:
+        return super().metadata_fields() + ["dimension"]
+
+    def _get_dimensions(self, value) -> List[int]:
+        dimensions: List[Value] = value["dimension"]
+        dims = [dim.size for dim in dimensions]
+        # Note: dimensions in NTNDArray are in fortran-like order
+        # with first index changing fastest.
+        #
+        # Therefore we need to reverse the order of the dimensions
+        # here to get back to a more usual C-like order with the
+        # last index changing fastest.
+        return dims[::-1]
+
+    def value(self, value):
+        dims = self._get_dimensions(value)
+        return value["value"].reshape(dims)
+
+    def descriptor(self, source: str, value) -> Descriptor:
+        dims = self._get_dimensions(value)
+        return dict(source=source, dtype="array", shape=dims)
+
+    def write_value(self, value):
+        # No clear use-case for writing directly to an NDArray, and some
+        # complexities around flattening to 1-D - e.g. dimension-order.
+        # Don't support this for now.
+        raise TypeError("Writing to NDArray not supported")
+
+
 @dataclass
 class PvaEnumConverter(PvaConverter):
     enum_class: Type[Enum]
@@ -122,7 +165,10 @@ def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConve
                 raise TypeError(f"{pv} has type [{pv_dtype}] not {datatype.__name__}")
             if dtype != pv_dtype:
                 raise TypeError(f"{pv} has type [{pv_dtype}] not [{dtype}]")
-        return PvaArrayConverter()
+        if "NTNDArray" in typeid:
+            return PvaNDArrayConverter()
+        else:
+            return PvaArrayConverter()
     elif "NTEnum" in typeid and datatype is bool:
         # Wanted a bool, but database represents as an enum
         pv_choices_len = get_unique(
@@ -213,14 +259,19 @@ async def get_descriptor(self) -> Descriptor:
         value = await self.ctxt.get(self.read_pv)
         return self.converter.descriptor(self.source, value)
 
+    def _pva_request_string(self, fields: List[str]) -> str:
+        return f"field({','.join(fields)})"
+
     async def get_reading(self) -> Reading:
-        value = await self.ctxt.get(
-            self.read_pv, request="field(value,alarm,timeStamp)"
+        request: str = self._pva_request_string(
+            self.converter.value_fields() + self.converter.metadata_fields()
         )
+        value = await self.ctxt.get(self.read_pv, request=request)
         return self.converter.reading(value)
 
     async def get_value(self) -> T:
-        value = await self.ctxt.get(self.read_pv, "field(value)")
+        request: str = self._pva_request_string(self.converter.value_fields())
+        value = await self.ctxt.get(self.read_pv, request)
         return self.converter.value(value)
 
     def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:

ophyd/v2/tests/test_epics.py

@@ -5,6 +5,7 @@
 import subprocess
 import sys
 import time
+from contextlib import closing
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -304,28 +305,27 @@ async def test_pva_ntdarray(ioc: IOC):
     if ioc.protocol == "ca":
         # CA can't do ndarray
         return
-    initial = np.zeros(4, np.int64)
-    put = np.ones_like(initial)
 
-    descriptor = dict(dtype="array", shape=[4])
+    put = np.array([1, 2, 3, 4, 5, 6], dtype=np.int64).reshape((2, 3))
+    initial = np.zeros_like(put)
 
+    backend = await ioc.make_backend(npt.NDArray[np.int64], "ntndarray")
+
+    # Backdoor into the "raw" data underlying the NDArray in QSrv
+    raw_data_backend = await ioc.make_backend(npt.NDArray[np.int64], "ntndarray:data")
+
+    # Make a monitor queue that will monitor for updates
     for i, p in [(initial, put), (put, initial)]:
-        backend = await ioc.make_backend(npt.NDArray[np.int64], "ntndarray")
-        # Make a monitor queue that will monitor for updates
-        q = MonitorQueue(backend)
-        try:
-            # Check descriptor
-            assert (
-                dict(source=backend.source, **descriptor)
-                == await backend.get_descriptor()
-            )
+        with closing(MonitorQueue(backend)) as q:
+            assert {
+                "source": backend.source,
+                "dtype": "array",
+                "shape": [2, 3],
+            } == await backend.get_descriptor()
             # Check initial value
             await q.assert_updates(pytest.approx(i))
-            # Put to new value and check that
-            await backend.put(p)
+            await raw_data_backend.put(p.flatten())
             await q.assert_updates(pytest.approx(p))
-        finally:
-            q.close()
 
 
 async def test_non_existant_errors(ioc: IOC):

ophyd/v2/tests/test_records.db

@@ -243,11 +243,31 @@ record(waveform, "$(P)table:enum")
     })
 }
 
-record(waveform, "$(P)ntndarray")
+record(longout, "$(P)ntndarray:ArraySize0_RBV") {
+    field(VAL, "3")
+    field(PINI, "YES")
+    info(Q:group, {
+        "$(P)ntndarray":{
+            "dimension[0].size":{+channel:"VAL", +type:"plain", +putorder:0}
+        }
+    })
+}
+
+record(longout, "$(P)ntndarray:ArraySize1_RBV") {
+    field(VAL, "2")
+    field(PINI, "YES")
+    info(Q:group, {
+        "$(P)ntndarray":{
+            "dimension[1].size":{+channel:"VAL", +type:"plain", +putorder:0}
+        }
+    })
+}
+
+record(waveform, "$(P)ntndarray:data")
 {
   field(FTVL, "INT64")
-  field(NELM, "4")
-  field(INP, {const:[0, 0, 0, 0]})
+  field(NELM, "6")
+  field(INP, {const:[0, 0, 0, 0, 0, 0]})
   field(PINI, "YES")
   info(Q:group, {
         "$(P)ntndarray":{
@@ -257,7 +277,7 @@ record(waveform, "$(P)ntndarray")
                 +channel:"VAL",
                 +trigger:"*",
             },
-            "": {+type:"meta", +channel:"VAL"}
+            "": {+type:"meta", +channel:"SEVR"}
         }
     })
 }