Remove apache-beam from dependencies to prevent Colab install error

datadrivenpdes/pipelines/__init__.py

@@ -1,3 +1,2 @@
 """Training and evaluation pipelines."""
-from datadrivenpdes.pipelines import beamlib
 from datadrivenpdes.pipelines import model_utils

datadrivenpdes/pipelines/beamlib.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 """Beam utilities."""
+
 import apache_beam as beam
 import numpy as np
 from typing import List, Tuple

datadrivenpdes/pipelines/create_training_data.py

@@ -19,12 +19,17 @@
 
 from absl import app
 from absl import flags
-import apache_beam as beam
+
+try:
+  import apache_beam as beam
+  from datadrivenpdes.pipelines import beamlib
+except:
+  pass
+
 import numpy as np
 from datadrivenpdes.core import builders
 from datadrivenpdes.core import equations
 from datadrivenpdes.core import grids
-from datadrivenpdes.pipelines import beamlib
 import tensorflow as tf
 
 # Ensure Equation subclasses are defined so we can look them up by name.

datadrivenpdes/pipelines/create_training_data_test.py

@@ -19,7 +19,12 @@
 from absl import flags
 from absl.testing import flagsaver
 from absl.testing import parameterized
-import apache_beam as beam
+
+try:
+  import apache_beam as beam
+except:
+  pass
+
 from datadrivenpdes.core import builders
 from datadrivenpdes.pipelines import create_training_data
 from tensorflow import gfile

setup.py

@@ -17,7 +17,6 @@
 
 INSTALL_REQUIRES = [
     'absl-py',
-    'apache-beam',
     'numpy',
     'scipy',
     'xarray',