<commit_before>import os
from flask import jsonify, request, url_for, make_response, abort
from flask_api import status
from .models.shopcart import Shopcart
from .models.product import Product
from .models.dataerror import DataValidationError
from . import app

######################################################################
# GET INDEX
######################################################################
@app.route('/')
def index():
    return app.send_static_file('index.html')

######################################################################
# RETRIEVE A SHOPCART
######################################################################
@app.route('/shopcarts/<int:id>', methods=['GET'])
def get_shopcarts(id):
    """
    Get a Shopcart
    This endpoint will get a shopcart.
    ---
    tags:
      - Shopcarts
    produces:
      - application/json
    parameters:
      - name: id
        in: path
        description: User id shopcart to retrieve
        type: integer
        required: true
    responses:
      200:
        description: Shopcart found
        schema:
          id: Shopcart
          properties:
            id:
              type: integer
              description: user id
              required: true
              default: 1
            products:
              type: dict
              description: products in the shopcart
              required: true
              default: {
                "2": {
                  "description": "Stationery",
                  "id": 2,
                  "name": "Pen",
                  "price": 3.4,
                  "quantity": 1
                }
              }
      404:
        description: Shopcart not found
    """
    cart = Shopcart.find(id)
    if cart:
        message = cart.serialize()
        rc = status.HTTP_200_OK
    else:
        message = {'error' : 'Shopcart with id: %s was not found' % str(id)}
        rc = status.HTTP_404_NOT_FOUND

    return jsonify(message), rc

######################################################################
# DELETE A SHOPCART
######################################################################
@app.route('/shopcarts/<int:id>', methods=['DELETE'])
def delete_shopcarts(id):
    """
    Delete a Shopcart
    This endpoint will delete a Shopcart based the id specified in the path
    ---
    tags:
      - Shopcarts
    description: Deletes a Shopcart from the database
    parameters:
      - name: id
        in: path
        description: ID of Shopcart to delete
        type: integer
        required: true
    responses:
      204:
        description: Shopcart deleted
    """
    cart = Shopcart.find(id)

    if cart:
        cart.delete()

    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# Create a Shopcart
######################################################################
@app.route('/shopcarts', methods=['POST'])
def create_shopcart():
    """
    Create a shopcart
    This endpoint will create a shopcart based in the data in the body that is posted.
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    parameters:
          - in: body
            user: body
            description: Data for the POST
            name: body
            schema:
              id: data
              required:
                - user_id
                - products
              properties:
                user_id:
                  type: integer
                  description: Unique id for the user
                products:
                  type: dict
                  description: list of dictionaries with products and quantities to add.
                  default: [{'pid':1, 'quantity':2},{'pid':3, 'quantity':4} ]
    responses:
      201:
        description: Shopcart created
        schema:
          properties:
            user_id:
              type: integer
              description: user id
              default: 1
            products:
              type: dict
              description: products in the shopcart
              required: true
              default: {
                "2": {
                  "description": "Stationery",
                  "id": 2,
                  "name": "Pen",
                  "price": 3.4,
                  "quantity": 1
                }
              }
      400:
        description: POST needs a valid user id
      409:
        description: Shopcart for that user already exits
    """
    data = request.get_json()
    try:
        user_id = data['user_id']
    except KeyError:
        message = {'error': 'POST needs a user id'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    cart = Shopcart.find(user_id)
    if cart:
        message = jsonify({'error' : 'Shopcart for user %s already exits' % str(user_id)})
        rc = status.HTTP_409_CONFLICT
        return make_response(message, rc)

    # Create the Cart
    products = data['products'] if 'products' in data else None

    try:
        if isinstance(products,dict):
            products = to_new_format(products)
        cart = Shopcart(user_id, products)
    except DataValidationError as e:
        message = {'error': e.args[0]}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    # If correct save it and return object
    cart.save()
    message = cart.serialize()
    location_url = url_for('get_shopcarts', id=int(cart.user_id), _external=True)
    return make_response(jsonify(message), status.HTTP_201_CREATED, {'Location': location_url})

######################################################################
# UPDATE AN EXISTING Shopcart product
######################################################################
@app.route('/shopcarts/<int:user_id>/products/<int:pid>', methods=['PUT'])
def update_shopcart(user_id, pid):
    """
    Update a Product in a given Shopcart
    This endpoint will update a Product in a given Shopcart based on the data in the body that is posted
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    description: The Shopcarts endpoint allows you to update a Product from a Shopcart
    parameters:
      - name: user_id
        in: path
        description: ID of Shopcart to retrieve
        type: integer
        required: true
      - name: pid
        in: path
        description: ID of product to update
        type: integer
        required: true
      - in: body
        name: body
        description: quantity to update
        type: body
        schema:
          id: dict
          properties:
            quantity:
              type: integer
          default: {"quantity": 10}
    responses:
      200:
        description: Product updated
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            product_id:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
      400:
        description: Bad Request (the posted data needs a quantity)
      404:
        description: Bad Request (the shopcart or product  was not found)
      """
    check_content_type('application/json')
    cart = Shopcart.find(user_id)

    if not cart:
        message = {'error' : 'Shopcart with id: %d was not found' % user_id}
        rc = status.HTTP_404_NOT_FOUND
        return make_response(jsonify(message), rc)

    prods_cart = cart.serialize()['products']
    if not pid in prods_cart.keys():
        message = {'error' : 'Product %d is not on shopcart %d' % (pid, user_id)}
        rc = status.HTTP_404_NOT_FOUND
        return make_response(jsonify(message), rc)

    data = request.get_json()
    try:
        quantity = data['quantity']
    except KeyError:
        message = {'error': 'Update product needs a quantity'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    if not isinstance(quantity, int):
        message = {'error': 'Update product needs a valid quantity'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    cart.update_product(pid, quantity)
    cart.save()

    return make_response(jsonify(cart.serialize()), status.HTTP_200_OK)

######################################################################
# ADD A PRODUCT TO A SHOPCART
######################################################################
@app.route('/shopcarts/<int:user_id>/products', methods=['POST'])
def add_product(user_id):
    """
    Add a Product to a Shopcart
    This endpoint will add a Product to a Shopcart based on the data in the body that is posted
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    description: The Shopcarts endpoint allows you to add a Product to a Shopcart
    parameters:
      - name: id
        in: path
        description: ID of Shopcart to retrieve
        type: integer
        required: true
      - in: body
        name: body
        schema:
          id: data
          required:
            - pid
            - quantity
          properties:
            pid:
              type: integer
              description: Unique id of Product to add
            quantity:
              type: integer
              description: Amount of Product to add
    responses:
      200:
        description: Product Added
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            product_id:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
      400:
        description: Bad Request (the posted data was not valid)
      """
    cart = Shopcart.find(user_id)

    if not cart:
        return jsonify("Cart with id '{}' was not found.".format(user_id)), status.HTTP_404_NOT_FOUND

    try:
        products = request.get_json()
        if isinstance(products,dict):
            products = to_new_format(products)
        cart.add_products(products)
        cart.save()
    except DataValidationError as e:
        message = {'error': e.args[0]}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    return make_response(jsonify(cart.serialize()), status.HTTP_200_OK)

######################################################################
# DELETE A PRODUCT FROM A SHOPCART
######################################################################
@app.route('/shopcarts/<int:user_id>/products/<int:pid>', methods=['DELETE'])
def delete_product(user_id, pid):
    cart = Shopcart.find(user_id)
    if cart:
        cart.delete_product(pid)
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# List all Shopcarts
######################################################################
@app.route('/shopcarts', methods=['GET'])
def get_all_shopcarts():

    """
    Retrieve a list of Shopcart
    This endpoint will return all Shopcarts unless a query parameter is specified
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to query Shopcarts
    parameters:
      - name: pid
        in: query
        description: the Product ID you want to query Shopcarts by
        required: false
        type: string
    responses:
      200:
        description: An array of Shopcarts
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            pid:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
                  description: Products in the Shopcart
      """

    pid = request.args.get('pid')
    if pid:
        carts = Shopcart.find_by_product( int(pid) )
    else:
        carts = Shopcart.all()
    message = [cart.serialize() for cart in carts]
    rc = status.HTTP_200_OK
    return jsonify(message), rc

######################################################################
# Prune empty Shopcarts
######################################################################
@app.route('/shopcarts/prune', methods=['DELETE'])
def prune_empty_shopcarts():
    """
    Prune all Shopcarts
    This endpoint will remove all empty Shopcarts (i.e. Shopcarts with no Products)
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to Prune all empty Shopcarts
    responses:
      204:
        description: Shopcarts have been pruned
      """
    Shopcart.prune()
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# Get all available Products
######################################################################
@app.route('/products', methods=['GET'])
def get_products():
    products = Product.all()
    message = [product.serialize() for product in products]

    return jsonify(message), status.HTTP_200_OK

######################################################################
# DELETE ALL SHOPCART DATA (for testing only)
######################################################################
@app.route('/shopcarts/reset', methods=['DELETE'])
def shopcarts_reset():
    """
    Reset  Shopcarts
    This endpoint will remove ALL Shopcarts
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to remove all Shopcarts
    responses:
      204:
        description: All Shopcarts have been removed
      """
    Shopcart.remove_all()
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
#  U T I L I T Y   F U N C T I O N S
######################################################################
def check_content_type(content_type):
    """ Checks that the media type is correct """
    if content_type in request.headers['Content-Type']:
        return
    #app.logger.error('Invalid Content-Type: %s', request.headers['Content-Type'])
    abort(status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, 'Content-Type must be {}'.format(content_type))

def init_db():
    """ Initializes the SQLAlchemy app """
    Shopcart.init_db()
    Product.seed_db()

def to_new_format(products):
    prods =[]
    if products:
        for pid,q in products.items():
            prods.append({'pid':pid,'quantity':q})
    return prods
<commit_msg>flassger delete a product<commit_after>import os
from flask import jsonify, request, url_for, make_response, abort
from flask_api import status
from .models.shopcart import Shopcart
from .models.product import Product
from .models.dataerror import DataValidationError
from . import app

######################################################################
# GET INDEX
######################################################################
@app.route('/')
def index():
    return app.send_static_file('index.html')

######################################################################
# RETRIEVE A SHOPCART
######################################################################
@app.route('/shopcarts/<int:id>', methods=['GET'])
def get_shopcarts(id):
    """
    Get a Shopcart
    This endpoint will get a shopcart.
    ---
    tags:
      - Shopcarts
    produces:
      - application/json
    parameters:
      - name: id
        in: path
        description: User id shopcart to retrieve
        type: integer
        required: true
    responses:
      200:
        description: Shopcart found
        schema:
          id: Shopcart
          properties:
            id:
              type: integer
              description: user id
              required: true
              default: 1
            products:
              type: dict
              description: products in the shopcart
              required: true
              default: {
                "2": {
                  "description": "Stationery",
                  "id": 2,
                  "name": "Pen",
                  "price": 3.4,
                  "quantity": 1
                }
              }
      404:
        description: Shopcart not found
    """
    cart = Shopcart.find(id)
    if cart:
        message = cart.serialize()
        rc = status.HTTP_200_OK
    else:
        message = {'error' : 'Shopcart with id: %s was not found' % str(id)}
        rc = status.HTTP_404_NOT_FOUND

    return jsonify(message), rc

######################################################################
# DELETE A SHOPCART
######################################################################
@app.route('/shopcarts/<int:id>', methods=['DELETE'])
def delete_shopcarts(id):
    """
    Delete a Shopcart
    This endpoint will delete a Shopcart based the id specified in the path
    ---
    tags:
      - Shopcarts
    description: Deletes a Shopcart from the database
    parameters:
      - name: id
        in: path
        description: ID of Shopcart to delete
        type: integer
        required: true
    responses:
      204:
        description: Shopcart deleted
    """
    cart = Shopcart.find(id)

    if cart:
        cart.delete()

    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# Create a Shopcart
######################################################################
@app.route('/shopcarts', methods=['POST'])
def create_shopcart():
    """
    Create a shopcart
    This endpoint will create a shopcart based in the data in the body that is posted.
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    parameters:
          - in: body
            user: body
            description: Data for the POST
            name: body
            schema:
              id: data
              required:
                - user_id
                - products
              properties:
                user_id:
                  type: integer
                  description: Unique id for the user
                products:
                  type: dict
                  description: list of dictionaries with products and quantities to add.
                  default: [{'pid':1, 'quantity':2},{'pid':3, 'quantity':4} ]
    responses:
      201:
        description: Shopcart created
        schema:
          properties:
            user_id:
              type: integer
              description: user id
              default: 1
            products:
              type: dict
              description: products in the shopcart
              required: true
              default: {
                "2": {
                  "description": "Stationery",
                  "id": 2,
                  "name": "Pen",
                  "price": 3.4,
                  "quantity": 1
                }
              }
      400:
        description: POST needs a valid user id
      409:
        description: Shopcart for that user already exits
    """
    data = request.get_json()
    try:
        user_id = data['user_id']
    except KeyError:
        message = {'error': 'POST needs a user id'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    cart = Shopcart.find(user_id)
    if cart:
        message = jsonify({'error' : 'Shopcart for user %s already exits' % str(user_id)})
        rc = status.HTTP_409_CONFLICT
        return make_response(message, rc)

    # Create the Cart
    products = data['products'] if 'products' in data else None

    try:
        if isinstance(products,dict):
            products = to_new_format(products)
        cart = Shopcart(user_id, products)
    except DataValidationError as e:
        message = {'error': e.args[0]}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    # If correct save it and return object
    cart.save()
    message = cart.serialize()
    location_url = url_for('get_shopcarts', id=int(cart.user_id), _external=True)
    return make_response(jsonify(message), status.HTTP_201_CREATED, {'Location': location_url})

######################################################################
# UPDATE AN EXISTING Shopcart product
######################################################################
@app.route('/shopcarts/<int:user_id>/products/<int:pid>', methods=['PUT'])
def update_shopcart(user_id, pid):
    """
    Update a Product in a given Shopcart
    This endpoint will update a Product in a given Shopcart based on the data in the body that is posted
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    description: The Shopcarts endpoint allows you to update a Product from a Shopcart
    parameters:
      - name: user_id
        in: path
        description: ID of Shopcart to retrieve
        type: integer
        required: true
      - name: pid
        in: path
        description: ID of product to update
        type: integer
        required: true
      - in: body
        name: body
        description: quantity to update
        type: body
        schema:
          id: dict
          properties:
            quantity:
              type: integer
          default: {"quantity": 10}
    responses:
      200:
        description: Product updated
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            product_id:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
      400:
        description: Bad Request (the posted data needs a quantity)
      404:
        description: Bad Request (the shopcart or product  was not found)
      """
    check_content_type('application/json')
    cart = Shopcart.find(user_id)

    if not cart:
        message = {'error' : 'Shopcart with id: %d was not found' % user_id}
        rc = status.HTTP_404_NOT_FOUND
        return make_response(jsonify(message), rc)

    prods_cart = cart.serialize()['products']
    if not pid in prods_cart.keys():
        message = {'error' : 'Product %d is not on shopcart %d' % (pid, user_id)}
        rc = status.HTTP_404_NOT_FOUND
        return make_response(jsonify(message), rc)

    data = request.get_json()
    try:
        quantity = data['quantity']
    except KeyError:
        message = {'error': 'Update product needs a quantity'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    if not isinstance(quantity, int):
        message = {'error': 'Update product needs a valid quantity'}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    cart.update_product(pid, quantity)
    cart.save()

    return make_response(jsonify(cart.serialize()), status.HTTP_200_OK)

######################################################################
# ADD A PRODUCT TO A SHOPCART
######################################################################
@app.route('/shopcarts/<int:user_id>/products', methods=['POST'])
def add_product(user_id):
    """
    Add a Product to a Shopcart
    This endpoint will add a Product to a Shopcart based on the data in the body that is posted
    ---
    tags:
      - Shopcarts
    consumes:
      - application/json
    produces:
      - application/json
    description: The Shopcarts endpoint allows you to add a Product to a Shopcart
    parameters:
      - name: user_id
        in: path
        description: ID of Shopcart to retrieve
        type: integer
        required: true
      - in: body
        name: body
        description: List of products/quantities
        type: array
        schema:
          id: data
          required:
            - pid
            - quantity
          properties:
            pid:
              type: integer
              description: Unique id of Product to add
            quantity:
              type: integer
              description: Amount of Product to add
          description: list of products/quantity to add
          default: {"pid": 1,"quantity": 10}
    responses:
      200:
        description: Product Added
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            product_id:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
      400:
        description: Bad Request (the posted data was not valid)
      """
    cart = Shopcart.find(user_id)

    if not cart:
        return jsonify("Cart with id '{}' was not found.".format(user_id)), status.HTTP_404_NOT_FOUND

    try:
        products = request.get_json()
        if isinstance(products,dict):
            products = to_new_format(products)
        cart.add_products(products)
        cart.save()
    except DataValidationError as e:
        message = {'error': e.args[0]}
        return jsonify(message), status.HTTP_400_BAD_REQUEST

    return make_response(jsonify(cart.serialize()), status.HTTP_200_OK)

######################################################################
# DELETE A PRODUCT FROM A SHOPCART
######################################################################
@app.route('/shopcarts/<int:user_id>/products/<int:pid>', methods=['DELETE'])
def delete_product(user_id, pid):
    """
    Delete a product from a Shopcart
    This endpoint will delete a product from a Shopcart
    based the user_id  and product_id specified in the path
    ---
    tags:
      - Shopcarts
    description: Delete a product from a Shopcart
    parameters:
      - name: user_id
        in: path
        description: ID of Shopcart
        type: integer
        required: true
      - name: pid
        in: path
        description: ID of the product
        type: integer
        required: true
    responses:
      204:
        description: Product deleted
    """
    cart = Shopcart.find(user_id)
    if cart:
        cart.delete_product(pid)
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# List all Shopcarts
######################################################################
@app.route('/shopcarts', methods=['GET'])
def get_all_shopcarts():

    """
    Retrieve a list of Shopcart
    This endpoint will return all Shopcarts unless a query parameter is specified
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to query Shopcarts
    parameters:
      - name: pid
        in: query
        description: the Product ID you want to query Shopcarts by
        required: false
        type: string
    responses:
      200:
        description: An array of Shopcarts
        schema:
          type: array
          items:
            schema:
              id: Shopcarts
              properties:
                user_id:
                  type: integer
                  description: Shopcart's unique ID associated with a user
                products:
                  type: array
                  items:
                    schema:
                        id: Products
                        properties:
                            pid:
                                type: integer
                                description: Product's unique id
                            name:
                                type: string
                                description: Name of the product
                            price:
                                type: integer
                                description: Cost of the product
                            description:
                                type: string
                                description: Description of the product
                  description: Products in the Shopcart
      """

    pid = request.args.get('pid')
    if pid:
        carts = Shopcart.find_by_product( int(pid) )
    else:
        carts = Shopcart.all()
    message = [cart.serialize() for cart in carts]
    rc = status.HTTP_200_OK
    return jsonify(message), rc

######################################################################
# Prune empty Shopcarts
######################################################################
@app.route('/shopcarts/prune', methods=['DELETE'])
def prune_empty_shopcarts():
    """
    Prune all Shopcarts
    This endpoint will remove all empty Shopcarts (i.e. Shopcarts with no Products)
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to Prune all empty Shopcarts
    responses:
      204:
        description: Shopcarts have been pruned
      """
    Shopcart.prune()
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
# Get all available Products
######################################################################
@app.route('/products', methods=['GET'])
def get_products():
    products = Product.all()
    message = [product.serialize() for product in products]

    return jsonify(message), status.HTTP_200_OK

######################################################################
# DELETE ALL SHOPCART DATA (for testing only)
######################################################################
@app.route('/shopcarts/reset', methods=['DELETE'])
def shopcarts_reset():
    """
    Reset  Shopcarts
    This endpoint will remove ALL Shopcarts
    ---
    tags:
      - Shopcarts
    description: The Shopcarts endpoint allows you to remove all Shopcarts
    responses:
      204:
        description: All Shopcarts have been removed
      """
    Shopcart.remove_all()
    return make_response('', status.HTTP_204_NO_CONTENT)

######################################################################
#  U T I L I T Y   F U N C T I O N S
######################################################################
def check_content_type(content_type):
    """ Checks that the media type is correct """
    if content_type in request.headers['Content-Type']:
        return
    #app.logger.error('Invalid Content-Type: %s', request.headers['Content-Type'])
    abort(status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, 'Content-Type must be {}'.format(content_type))

def init_db():
    """ Initializes the SQLAlchemy app """
    Shopcart.init_db()
    Product.seed_db()

def to_new_format(products):
    prods =[]
    if products:
        for pid,q in products.items():
            prods.append({'pid':pid,'quantity':q})
    return prods
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
Ssh-key key handlers and associated routines
"""
import contextlib
import logging
import paramiko
import re
from datetime import datetime

from cStringIO import StringIO
from teuthology import contextutil
import teuthology.misc as misc
from ..orchestra import run

log = logging.getLogger(__name__)
ssh_keys_user = 'ssh-keys-user'

def timestamp(format_='%Y-%m-%d_%H:%M:%S:%f'):
    """
    Return a UTC timestamp suitable for use in filenames
    """
    return datetime.utcnow().strftime(format_)

def backup_file(remote, path, sudo=False):
    """
    Creates a backup of a file on the remote, simply by copying it and adding a
    timestamp to the name.
    """
    backup_path = "{path}_{timestamp}".format(
        path=path, timestamp=timestamp()
    )
    args = [
        'cp', '-v', '-a', path, backup_path,
    ]
    if sudo:
        args.insert(0, 'sudo')
    remote.run(args=args)
    return backup_path

def generate_keys():
    """
    Generatees a public and private key
    """
    key = paramiko.RSAKey.generate(2048)
    privateString = StringIO()
    key.write_private_key(privateString)
    return key.get_base64(), privateString.getvalue()

def particular_ssh_key_test(line_to_test, ssh_key):
    """
    Check the validity of the ssh_key
    """
    match = re.match('[\w-]+ {key} \S+@\S+'.format(key=re.escape(ssh_key)), line_to_test)

    if match:
        return False
    else:
        return True

def ssh_keys_user_line_test(line_to_test, username ):
    """
    Check the validity of the username
    """
    match = re.match('[\w-]+ \S+ {username}@\S+'.format(username=username), line_to_test)

    if match:
        return False
    else:
        return True

def cleanup_added_key(ctx, orig_path, path):
    """
    Delete the keys and removes ~/.ssh/authorized_keys entries we added
    """
    log.info('cleaning up keys added for testing')

    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('  cleaning up keys for user {user} on {host}'.format(host=hostname, user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa'.format(user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa.pub'.format(user=username))
            misc.move_file(remote, orig_path, path)

@contextlib.contextmanager
def tweak_ssh_config(ctx, config):
    """
    Turn off StrictHostKeyChecking
    """
    run.wait(
        ctx.cluster.run(
            args=[
                'echo',
                'StrictHostKeyChecking no\n',
                run.Raw('>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                'echo',
                'UserKnownHostsFile ',
                run.Raw('/dev/null'),
                run.Raw('>>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                run.Raw('chmod 600 /home/ubuntu/.ssh/config'),
            ],
            wait=False,
        )
    )

    try:
        yield

    finally:
        run.wait(
            ctx.cluster.run(
                args=['rm',run.Raw('/home/ubuntu/.ssh/config')],
            wait=False
            ),
        )

@contextlib.contextmanager
def push_keys_to_host(ctx, config, public_key, private_key):
    """
    Push keys to all hosts
    """
    log.info('generated public key {pub_key}'.format(pub_key=public_key))

    # add an entry for all hosts in ctx to auth_keys_data
    auth_keys_data = ''

    for inner_host in ctx.cluster.remotes.iterkeys():
        inner_username, inner_hostname = str(inner_host).split('@')
        # create a 'user@hostname' string using our fake hostname
        fake_hostname = '{user}@{host}'.format(user=ssh_keys_user, host=str(inner_hostname))
        auth_keys_data += '\nssh-rsa {pub_key} {user_host}\n'.format(pub_key=public_key, user_host=fake_hostname)

    # for each host in ctx, add keys for all other hosts
    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('pushing keys to {host} for {user}'.format(host=hostname, user=username))

            # adding a private key
            priv_key_file = '/home/{user}/.ssh/id_rsa'.format(user=username)
            priv_key_data = '{priv_key}'.format(priv_key=private_key)
            misc.delete_file(remote, priv_key_file, force=True)
            # Hadoop requires that .ssh/id_rsa have permissions of '500'
            misc.create_file(remote, priv_key_file, priv_key_data, str(500))

            # then a private key
            pub_key_file = '/home/{user}/.ssh/id_rsa.pub'.format(user=username)
            pub_key_data = 'ssh-rsa {pub_key} {user_host}'.format(pub_key=public_key, user_host=str(remote))
            misc.delete_file(remote, pub_key_file, force=True)
            misc.create_file(remote, pub_key_file, pub_key_data)

            # add appropriate entries to the authorized_keys file for this host
            auth_keys_file = '/home/{user}/.ssh/authorized_keys'.format(
                user=username)
            orig_auth_keys_file = backup_file(remote, auth_keys_file)
            misc.append_lines_to_file(remote, auth_keys_file, auth_keys_data)

    try:
        yield

    finally:
        # cleanup the keys
        log.info("Cleaning up SSH keys")
        cleanup_added_key(ctx, orig_auth_keys_file, auth_keys_file)

@contextlib.contextmanager
def task(ctx, config):
    """
    Creates a set of RSA keys, distributes the same key pair
    to all hosts listed in ctx.cluster, and adds all hosts
    to all others authorized_keys list.

    During cleanup it will delete .ssh/id_rsa, .ssh/id_rsa.pub
    and remove the entries in .ssh/authorized_keys while leaving
    pre-existing entries in place.
    """

    if config is None:
        config = {}
    assert isinstance(config, dict), \
        "task hadoop only supports a dictionary for configuration"

    # this does not need to do cleanup and does not depend on
    # ctx, so I'm keeping it outside of the nested calls
    public_key_string, private_key_string = generate_keys()

    with contextutil.nested(
        lambda: tweak_ssh_config(ctx, config),
        lambda: push_keys_to_host(ctx, config, public_key_string, private_key_string),
        #lambda: tweak_ssh_config(ctx, config),
        ):
        yield

<commit_msg>ssh_keys: save each remote's backup file name separately in a dict for undo<commit_after>#!/usr/bin/python
"""
Ssh-key key handlers and associated routines
"""
import contextlib
import logging
import paramiko
import re
from datetime import datetime

from cStringIO import StringIO
from teuthology import contextutil
import teuthology.misc as misc
from ..orchestra import run

log = logging.getLogger(__name__)
ssh_keys_user = 'ssh-keys-user'

def timestamp(format_='%Y-%m-%d_%H:%M:%S:%f'):
    """
    Return a UTC timestamp suitable for use in filenames
    """
    return datetime.utcnow().strftime(format_)

def backup_file(remote, path, sudo=False):
    """
    Creates a backup of a file on the remote, simply by copying it and adding a
    timestamp to the name.
    """
    backup_path = "{path}_{timestamp}".format(
        path=path, timestamp=timestamp()
    )
    args = [
        'cp', '-v', '-a', path, backup_path,
    ]
    if sudo:
        args.insert(0, 'sudo')
    remote.run(args=args)
    return backup_path

def generate_keys():
    """
    Generatees a public and private key
    """
    key = paramiko.RSAKey.generate(2048)
    privateString = StringIO()
    key.write_private_key(privateString)
    return key.get_base64(), privateString.getvalue()

def particular_ssh_key_test(line_to_test, ssh_key):
    """
    Check the validity of the ssh_key
    """
    match = re.match('[\w-]+ {key} \S+@\S+'.format(key=re.escape(ssh_key)), line_to_test)

    if match:
        return False
    else:
        return True

def ssh_keys_user_line_test(line_to_test, username ):
    """
    Check the validity of the username
    """
    match = re.match('[\w-]+ \S+ {username}@\S+'.format(username=username), line_to_test)

    if match:
        return False
    else:
        return True

def cleanup_added_key(ctx, key_backup_files, path):
    """
    Delete the keys and removes ~/.ssh/authorized_keys entries we added
    """
    log.info('cleaning up keys added for testing')

    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('  cleaning up keys for user {user} on {host}'.format(host=hostname, user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa'.format(user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa.pub'.format(user=username))
            misc.move_file(remote, key_backup_files[remote], path)

@contextlib.contextmanager
def tweak_ssh_config(ctx, config):
    """
    Turn off StrictHostKeyChecking
    """
    run.wait(
        ctx.cluster.run(
            args=[
                'echo',
                'StrictHostKeyChecking no\n',
                run.Raw('>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                'echo',
                'UserKnownHostsFile ',
                run.Raw('/dev/null'),
                run.Raw('>>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                run.Raw('chmod 600 /home/ubuntu/.ssh/config'),
            ],
            wait=False,
        )
    )

    try:
        yield

    finally:
        run.wait(
            ctx.cluster.run(
                args=['rm',run.Raw('/home/ubuntu/.ssh/config')],
            wait=False
            ),
        )

@contextlib.contextmanager
def push_keys_to_host(ctx, config, public_key, private_key):
    """
    Push keys to all hosts
    """
    log.info('generated public key {pub_key}'.format(pub_key=public_key))

    # add an entry for all hosts in ctx to auth_keys_data
    auth_keys_data = ''

    for inner_host in ctx.cluster.remotes.iterkeys():
        inner_username, inner_hostname = str(inner_host).split('@')
        # create a 'user@hostname' string using our fake hostname
        fake_hostname = '{user}@{host}'.format(user=ssh_keys_user, host=str(inner_hostname))
        auth_keys_data += '\nssh-rsa {pub_key} {user_host}\n'.format(pub_key=public_key, user_host=fake_hostname)

    key_backup_files = dict()
    # for each host in ctx, add keys for all other hosts
    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('pushing keys to {host} for {user}'.format(host=hostname, user=username))

            # adding a private key
            priv_key_file = '/home/{user}/.ssh/id_rsa'.format(user=username)
            priv_key_data = '{priv_key}'.format(priv_key=private_key)
            misc.delete_file(remote, priv_key_file, force=True)
            # Hadoop requires that .ssh/id_rsa have permissions of '500'
            misc.create_file(remote, priv_key_file, priv_key_data, str(500))

            # then a private key
            pub_key_file = '/home/{user}/.ssh/id_rsa.pub'.format(user=username)
            pub_key_data = 'ssh-rsa {pub_key} {user_host}'.format(pub_key=public_key, user_host=str(remote))
            misc.delete_file(remote, pub_key_file, force=True)
            misc.create_file(remote, pub_key_file, pub_key_data)

            # add appropriate entries to the authorized_keys file for this host
            auth_keys_file = '/home/{user}/.ssh/authorized_keys'.format(
                user=username)
            key_backup_files[remote] = backup_file(remote, auth_keys_file)
            misc.append_lines_to_file(remote, auth_keys_file, auth_keys_data)

    try:
        yield

    finally:
        # cleanup the keys
        log.info("Cleaning up SSH keys")
        cleanup_added_key(ctx, key_backup_files, auth_keys_file)

@contextlib.contextmanager
def task(ctx, config):
    """
    Creates a set of RSA keys, distributes the same key pair
    to all hosts listed in ctx.cluster, and adds all hosts
    to all others authorized_keys list.

    During cleanup it will delete .ssh/id_rsa, .ssh/id_rsa.pub
    and remove the entries in .ssh/authorized_keys while leaving
    pre-existing entries in place.
    """

    if config is None:
        config = {}
    assert isinstance(config, dict), \
        "task hadoop only supports a dictionary for configuration"

    # this does not need to do cleanup and does not depend on
    # ctx, so I'm keeping it outside of the nested calls
    public_key_string, private_key_string = generate_keys()

    with contextutil.nested(
        lambda: tweak_ssh_config(ctx, config),
        lambda: push_keys_to_host(ctx, config, public_key_string, private_key_string),
        #lambda: tweak_ssh_config(ctx, config),
        ):
        yield

<|endoftext|>


<commit_before># -----------------------------------------------------------------------------
# Copyright (c) 2014--, The Qiita Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------

from os.path import join, exists
from functools import partial
from os import makedirs

import pandas as pd
from h5py import File
# Currently the EBI submission is not part of the target gene plugin
# and the demux file is not its own project, so we need to import from
# qiita_ware. Plans are to the EBI submission to the target gene plugin
from qiita_ware.demux import to_hdf5

def get_artifact_information(qclient, artifact_id):
    """Retrieves the artifact information for running split libraries

    Parameters
    ----------
    qclient : tgp.qiita_client.QiitaClient
        The Qiita server client
    artifact_id : str
        The artifact id

    Returns
    -------
    list of (str, str), str, str
        The artifact filepaths and their type
        The artifact mapping file
        The artifact type

    Raises
    ------
    ValueError
        If there is any problem gathering the information from the server
    """
    # Get the artifact filepath information
    fps_info = qclient.get("/qiita_db/artifacts/%s/filepaths/" % artifact_id)
    if not fps_info or not fps_info['success']:
        error_msg = "Could not get artifact filepath information: %s"
        if fps_info:
            error_msg = error_msg % fps_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    fps = fps_info['filepaths']

    # Get the artifact metadata
    metadata_info = qclient.get(
        "/qiita_db/artifacts/%s/mapping/" % artifact_id)
    if not metadata_info or not metadata_info['success']:
        error_msg = "Could not get artifact metadata information %s"
        if metadata_info:
            error_msg = error_msg % metadata_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    mapping_file = metadata_info['mapping']

    # Get the artifact type
    type_info = qclient.get("/qiita_db/artifacts/%s/type/" % artifact_id)
    if not type_info or not type_info['success']:
        error_msg = "Could not get artifact metadata information %s"
        if type_info:
            error_msg = error_msg % type_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    artifact_type = type_info['type']

    return fps, mapping_file, artifact_type

def split_mapping_file(mapping_file, out_dir):
    """Splits a QIIME-compliant mapping file by run_prefix

    Parameters
    ----------
    mapping_file : str
        The mapping file filepath
    out_dir : str
        The path to the output directory

    Returns
    -------
    list of str
        The paths to the splitted mapping files
    """
    mf = pd.read_csv(mapping_file, delimiter='\t', dtype=str, encoding='utf-8')
    mf.set_index('#SampleID', inplace=True)

    path_builder = partial(join, out_dir)
    if 'run_prefix' in mf:
        if not exists(out_dir):
            makedirs(out_dir)
        output_fps = []
        for prefix, df in mf.groupby('run_prefix'):
            out_fp = path_builder('%s_mapping_file.txt' % prefix)
            output_fps.append(out_fp)
            df.to_csv(out_fp, index_label='#SampleID', sep='\t')
    else:
        output_fps = [mapping_file]

    return output_fps

def generate_demux_file(sl_out):
    """Creates the HDF5 demultiplexed file

    Parameters
    ----------
    sl_out : str
        Path to the output directory of split libraries

    Returns
    -------
    str
        The path of the demux file

    Raises
    ------
    ValueError
        If the split libraries output does not contain the demultiplexed fastq
        file
    """
    fastq_fp = str(join(sl_out, 'seqs.fastq'))
    if not exists(fastq_fp):
        raise ValueError("The split libraries output directory does not "
                         "contain the demultiplexed fastq file.")

    demux_fp = join(sl_out, 'seqs.demux')
    with File(demux_fp, "w") as f:
        to_hdf5(fastq_fp, f)
    return demux_fp

def generate_artifact_info(sl_out):
    """Creates the artifact information to attach to the payload

    Parameters
    ----------
    sl_out : str
        Path to the split libraries output directory

    Returns
    -------
    list of [str, list of (str, str), bool, bool]
        The artifacts information to include in the payload when the split
        libraries job is completed.
        - The artifact type
        - The list of filepaths with their artifact type
        - Whether the artifact can be submitted to ebi
        - Whether the artifact can be submitted to vamps
    """
    path_builder = partial(join, sl_out)
    filepaths = [(path_builder('seqs.fna'), 'preprocessed_fasta'),
                 (path_builder('seqs.fastq'), 'preprocessed_fastq'),
                 (path_builder('seqs.demux'), 'preprocessed_demux'),
                 (path_builder('split_library_log.txt'), 'log')]
    return [['Demultiplexed', filepaths, True, True]]
<commit_msg>Fixing split mapping file serializer<commit_after># -----------------------------------------------------------------------------
# Copyright (c) 2014--, The Qiita Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------

from os.path import join, exists
from functools import partial
from os import makedirs

import pandas as pd
from h5py import File
# Currently the EBI submission is not part of the target gene plugin
# and the demux file is not its own project, so we need to import from
# qiita_ware. Plans are to the EBI submission to the target gene plugin
from qiita_ware.demux import to_hdf5

def get_artifact_information(qclient, artifact_id):
    """Retrieves the artifact information for running split libraries

    Parameters
    ----------
    qclient : tgp.qiita_client.QiitaClient
        The Qiita server client
    artifact_id : str
        The artifact id

    Returns
    -------
    list of (str, str), str, str
        The artifact filepaths and their type
        The artifact mapping file
        The artifact type

    Raises
    ------
    ValueError
        If there is any problem gathering the information from the server
    """
    # Get the artifact filepath information
    fps_info = qclient.get("/qiita_db/artifacts/%s/filepaths/" % artifact_id)
    if not fps_info or not fps_info['success']:
        error_msg = "Could not get artifact filepath information: %s"
        if fps_info:
            error_msg = error_msg % fps_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    fps = fps_info['filepaths']

    # Get the artifact metadata
    metadata_info = qclient.get(
        "/qiita_db/artifacts/%s/mapping/" % artifact_id)
    if not metadata_info or not metadata_info['success']:
        error_msg = "Could not get artifact metadata information %s"
        if metadata_info:
            error_msg = error_msg % metadata_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    mapping_file = metadata_info['mapping']

    # Get the artifact type
    type_info = qclient.get("/qiita_db/artifacts/%s/type/" % artifact_id)
    if not type_info or not type_info['success']:
        error_msg = "Could not get artifact metadata information %s"
        if type_info:
            error_msg = error_msg % type_info['error']
        else:
            error_msg = error_msg % "could not connect with the server"
        raise ValueError(error_msg)
    artifact_type = type_info['type']

    return fps, mapping_file, artifact_type

def split_mapping_file(mapping_file, out_dir):
    """Splits a QIIME-compliant mapping file by run_prefix

    Parameters
    ----------
    mapping_file : str
        The mapping file filepath
    out_dir : str
        The path to the output directory

    Returns
    -------
    list of str
        The paths to the splitted mapping files
    """
    mf = pd.read_csv(mapping_file, delimiter='\t', dtype=str, encoding='utf-8')
    mf.set_index('#SampleID', inplace=True)

    path_builder = partial(join, out_dir)
    if 'run_prefix' in mf:
        if not exists(out_dir):
            makedirs(out_dir)
        output_fps = []
        for prefix, df in mf.groupby('run_prefix'):
            out_fp = path_builder('%s_mapping_file.txt' % prefix)
            output_fps.append(out_fp)
            df.to_csv(out_fp, index_label='#SampleID', sep='\t',
                      encoding='utf-8')
    else:
        output_fps = [mapping_file]

    return output_fps

def generate_demux_file(sl_out):
    """Creates the HDF5 demultiplexed file

    Parameters
    ----------
    sl_out : str
        Path to the output directory of split libraries

    Returns
    -------
    str
        The path of the demux file

    Raises
    ------
    ValueError
        If the split libraries output does not contain the demultiplexed fastq
        file
    """
    fastq_fp = str(join(sl_out, 'seqs.fastq'))
    if not exists(fastq_fp):
        raise ValueError("The split libraries output directory does not "
                         "contain the demultiplexed fastq file.")

    demux_fp = join(sl_out, 'seqs.demux')
    with File(demux_fp, "w") as f:
        to_hdf5(fastq_fp, f)
    return demux_fp

def generate_artifact_info(sl_out):
    """Creates the artifact information to attach to the payload

    Parameters
    ----------
    sl_out : str
        Path to the split libraries output directory

    Returns
    -------
    list of [str, list of (str, str), bool, bool]
        The artifacts information to include in the payload when the split
        libraries job is completed.
        - The artifact type
        - The list of filepaths with their artifact type
        - Whether the artifact can be submitted to ebi
        - Whether the artifact can be submitted to vamps
    """
    path_builder = partial(join, sl_out)
    filepaths = [(path_builder('seqs.fna'), 'preprocessed_fasta'),
                 (path_builder('seqs.fastq'), 'preprocessed_fastq'),
                 (path_builder('seqs.demux'), 'preprocessed_demux'),
                 (path_builder('split_library_log.txt'), 'log')]
    return [['Demultiplexed', filepaths, True, True]]
<|endoftext|>


<commit_before># Standard library
from collections import OrderedDict
import copy

# Third-party
import numpy as np

# Package
from ..utils import quantity_to_hdf5, quantity_from_hdf5

__all__ = ['JokerSamples']

class JokerSamples(OrderedDict):

    def __init__(self, trend_cls=None, **kwargs):
        """ """

        self._valid_keys = ['P', 'phi0', 'ecc', 'omega', 'jitter', 'K']

        self.trend_cls = trend_cls
        if trend_cls is not None:
            self._valid_keys += trend_cls.parameters

        kw = kwargs.copy()

        self._n_samples = None
        for key,val in kw.items():
            self._validate_key(key)
            kw[key] = self._validate_val(val)

            if self._n_samples is None:
                self._n_samples = len(val)

        super(JokerSamples, self).__init__(**kw)

    def _validate_key(self, key):
        if key not in self._valid_keys:
            raise ValueError("Invalid key '{0}'.".format(key))

    def _validate_val(self, val):
        val = np.atleast_1d(val)
        if self._n_samples is not None and len(val) != self._n_samples:
            raise ValueError("Length of new samples must match those already "
                             "stored! ({0}, expected {1})"
                             .format(len(val), self._n_samples))

        return val

    def __getitem__(self, slc):
        if isinstance(slc, str):
            return super(JokerSamples, self).__getitem__(slc)

        else:
            new = copy.copy(self)
            new._n_samples = None # reset number of samples

            for k in self.keys():
                new[k] = self[k][slc]

            return new

    def __setitem__(self, key, val):
        self._validate_key(key)
        val = self._validate_val(val)

        if self._n_samples is None:
            self._n_samples = len(val)

        super(JokerSamples, self).__setitem__(key, val)

    @property
    def n_samples(self):
        if self._n_samples is None:
            raise ValueError("No samples stored!")
        return self._n_samples

    def __len__(self):
        return self.n_samples

    def __str__(self):
        return "<JokerSamples: n={}>".format()

    @classmethod
    def from_hdf5(cls, f, n=None):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        n : int (optional)
            The number of samples to load.
        """
        samples = cls()
        for key in f.keys():
            samples[key] = quantity_from_hdf5(f, key, n=n)

        return samples

    def to_hdf5(self, f):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        """

        for key in self.keys():
            quantity_to_hdf5(f, key, self[key])
<commit_msg>option to pass trend class in from_hdf5<commit_after># Standard library
from collections import OrderedDict
import copy

# Third-party
import numpy as np

# Package
from ..utils import quantity_to_hdf5, quantity_from_hdf5

__all__ = ['JokerSamples']

class JokerSamples(OrderedDict):

    def __init__(self, trend_cls=None, **kwargs):
        """ """

        self._valid_keys = ['P', 'phi0', 'ecc', 'omega', 'jitter', 'K']

        self.trend_cls = trend_cls
        if trend_cls is not None:
            self._valid_keys += trend_cls.parameters

        kw = kwargs.copy()

        self._n_samples = None
        for key,val in kw.items():
            self._validate_key(key)
            kw[key] = self._validate_val(val)

            if self._n_samples is None:
                self._n_samples = len(val)

        super(JokerSamples, self).__init__(**kw)

    def _validate_key(self, key):
        if key not in self._valid_keys:
            raise ValueError("Invalid key '{0}'.".format(key))

    def _validate_val(self, val):
        val = np.atleast_1d(val)
        if self._n_samples is not None and len(val) != self._n_samples:
            raise ValueError("Length of new samples must match those already "
                             "stored! ({0}, expected {1})"
                             .format(len(val), self._n_samples))

        return val

    def __getitem__(self, slc):
        if isinstance(slc, str):
            return super(JokerSamples, self).__getitem__(slc)

        else:
            new = copy.copy(self)
            new._n_samples = None # reset number of samples

            for k in self.keys():
                new[k] = self[k][slc]

            return new

    def __setitem__(self, key, val):
        self._validate_key(key)
        val = self._validate_val(val)

        if self._n_samples is None:
            self._n_samples = len(val)

        super(JokerSamples, self).__setitem__(key, val)

    @property
    def n_samples(self):
        if self._n_samples is None:
            raise ValueError("No samples stored!")
        return self._n_samples

    def __len__(self):
        return self.n_samples

    def __str__(self):
        return "<JokerSamples: n={}>".format()

    @classmethod
    def from_hdf5(cls, f, n=None, **kwargs):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        n : int (optional)
            The number of samples to load.
        **kwargs
            All other keyword arguments are passed to the class initializer.
            You may need to use this to pass in a ``trend_cls``.
        """
        samples = cls(**kwargs)
        for key in f.keys():
            samples[key] = quantity_from_hdf5(f, key, n=n)

        return samples

    def to_hdf5(self, f):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        """

        for key in self.keys():
            quantity_to_hdf5(f, key, self[key])
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2011  Renato Lima - Akretion                                  #
#                                                                             #
#This program is free software: you can redistribute it and/or modify         #
#it under the terms of the GNU Affero General Public License as published by  #
#the Free Software Foundation, either version 3 of the License, or            #
#(at your option) any later version.                                          #
#                                                                             #
#This program is distributed in the hope that it will be useful,              #
#but WITHOUT ANY WARRANTY; without even the implied warranty of               #
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the                #
#GNU Affero General Public License for more details.                          #
#                                                                             #
#You should have received a copy of the GNU Affero General Public License     #
#along with this program.  If not, see <http://www.gnu.org/licenses/>.        #
###############################################################################

from osv import fields, osv

class account_move_line(osv.Model):
    _inherit = 'account.move.line'
    _columns = {
        'revenue_expense': fields.related(
            'journal_id', 'revenue_expense', type='many2one',
            relation='account.journal', string='Revenue Expense',
            store=True, readonly=True),
    }
<commit_msg>Fixed small bug with the type of revenue_expense in account_payment_extension<commit_after># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2011  Renato Lima - Akretion                                  #
#                                                                             #
#This program is free software: you can redistribute it and/or modify         #
#it under the terms of the GNU Affero General Public License as published by  #
#the Free Software Foundation, either version 3 of the License, or            #
#(at your option) any later version.                                          #
#                                                                             #
#This program is distributed in the hope that it will be useful,              #
#but WITHOUT ANY WARRANTY; without even the implied warranty of               #
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the                #
#GNU Affero General Public License for more details.                          #
#                                                                             #
#You should have received a copy of the GNU Affero General Public License     #
#along with this program.  If not, see <http://www.gnu.org/licenses/>.        #
###############################################################################

from osv import fields, osv

class account_move_line(osv.Model):
    _inherit = 'account.move.line'
    _columns = {
        'revenue_expense': fields.related(
            'journal_id', 'revenue_expense', type='boolean',
            string='Revenue Expense',
            store=True, readonly=True),
    }
<|endoftext|>


<commit_before>"""
Credit for demo setup: Gusto development team
"""
from gusto import *
from firedrake import as_vector, SpatialCoordinate,\
    PeriodicRectangleMesh, ExtrudedMesh, \
    exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt, \
    parameters
from firedrake.petsc import PETSc
from argparse import ArgumentParser
from hybridization import HybridizedCompressibleSolver
import sys

day = 24.*60.*60.
hour = 60.*60.
dt = 30.

PETSc.Log.begin()
parser = ArgumentParser(description="""Euler-Boussinesq Eady slice model.""",
                        add_help=False)

parser.add_argument("--hybridization",
                    action="store_true",
                    help="Use a hybridized compressible solver.")

parser.add_argument("--test",
                    action="store_true",
                    help="Enable a quick test run.")

parser.add_argument("--profile",
                    action="store_true",
                    help="Turn on profiling for a 20 time-step run.")

parser.add_argument("--help",
                    action="store_true",
                    help="Show help.")

args, _ = parser.parse_known_args()

if args.help:
    help = parser.format_help()
    PETSc.Sys.Print("%s\n" % help)
    sys.exit(1)

if args.profile:
    # Ensures accurate timing of parallel loops
    parameters["pyop2_options"]["lazy_evaluation"] = False
    tmax = 20*dt

if args.test:
    tmax = dt

if not args.test and not args.profile:
    tmax = 30*day

hybrid = bool(args.hybridization)
PETSc.Sys.Print("""
Problem parameters:\n
Test case: Euler-Boussinesq Eady model.\n
Hybridized compressible solver: %s,\n
Profiling: %s,\n
Test run: %s.\n
""" % (hybrid, bool(args.profile), bool(args.test)))

PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" % (dt,
                                                                      tmax))

# Construct 1d periodic base mesh
columns = 30  # number of columns
L = 1000000.
m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True)

# Build 2D mesh by extruding the base mesh
nlayers = 30  # horizontal layers
H = 10000.  # height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)

# Set up all the other things that state requires:
# Coriolis expression
f = 1.e-04
Omega = as_vector([0., 0., f*0.5])

# List of prognostic fieldnames:
# this is passed to state and used to construct a dictionary,
# state.field_dict so that we can access fields by name.
# u is the 3D velocity
# p is the pressure
# b is the buoyancy
fieldlist = ['u', 'rho', 'theta']

# Class containing timestepping parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
timestepping = TimesteppingParameters(dt=dt)

# Class containing output parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
if hybrid:
    dirname = "hybrid_compressible_eady"
else:
    dirname = "compressible_eady"

output = OutputParameters(dirname=dirname,
                          dumpfreq=int(2*hour/dt),
                          dumplist=['u', 'rho', 'theta'],
                          perturbation_fields=['rho', 'theta', 'ExnerPi'])

# Class containing physical parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
parameters = CompressibleEadyParameters(H=H, f=f)

# Class for diagnostics:
# fields passed to this class will have basic diagnostics computed
# (eg min, max, l2 norm) and these will be output as a json file
diagnostics = Diagnostics(*fieldlist)

# List of diagnostic fields, each defined in a class in diagnostics.py
diagnostic_fields = [CourantNumber(), VelocityY(),
                     ExnerPi(), ExnerPi(reference=True),
                     CompressibleKineticEnergy(),
                     CompressibleKineticEnergyY(),
                     CompressibleEadyPotentialEnergy(),
                     Sum("CompressibleKineticEnergy",
                         "CompressibleEadyPotentialEnergy"),
                     Difference("CompressibleKineticEnergy",
                                "CompressibleKineticEnergyY")]

# Setup state, passing in the mesh, information on the required finite element
# function spaces and the classes above
state = State(mesh, vertical_degree=1, horizontal_degree=1,
              family="RTCF",
              Coriolis=Omega,
              timestepping=timestepping,
              output=output,
              parameters=parameters,
              diagnostics=diagnostics,
              fieldlist=fieldlist,
              diagnostic_fields=diagnostic_fields)

# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")

# Spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()

# First setup the background buoyancy profile
# z.grad(bref) = N**2.
# The following is symbolic algebra, using the default buoyancy frequency
# from the parameters class.
x, y, z = SpatialCoordinate(mesh)
g = parameters.g
Nsq = parameters.Nsq
theta_surf = parameters.theta_surf

# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
theta_ref = theta_surf*exp(Nsq*(z-H/2)/g)
theta_b = Function(Vt).interpolate(theta_ref)

# Set theta_pert
def coth(x):
    return cosh(x)/sinh(x)

def Z(z):
    return Bu*((z/H)-0.5)

def n():
    return Bu**(-1)*sqrt((Bu*0.5-tanh(Bu*0.5))*(coth(Bu*0.5)-Bu*0.5))

a = -4.5
Bu = 0.5
theta_exp = a*theta_surf/g*sqrt(Nsq)*(-(1.-Bu*0.5*coth(Bu*0.5))*sinh(Z(z))*cos(pi*(x-L)/L)
                                      - n()*Bu*cosh(Z(z))*sin(pi*(x-L)/L))
theta_pert = Function(Vt).interpolate(theta_exp)

# Set theta0
theta0.interpolate(theta_b + theta_pert)

# Calculate hydrostatic Pi
PETSc.Sys.Print("Computing hydrostatic Exner pressure...\n")
rho_b = Function(Vr)
compressible_hydrostatic_balance(state, theta_b, rho_b)
compressible_hydrostatic_balance(state, theta0, rho0)

# Set Pi0
Pi0 = calculate_Pi0(state, theta0, rho0)
state.parameters.Pi0 = Pi0

# Set x component of velocity
cp = state.parameters.cp
dthetady = state.parameters.dthetady
Pi = exner(theta0, rho0, state)
u = cp*dthetady/f*(Pi-Pi0)

# Set y component of velocity
v = Function(Vr).assign(0.)
compressible_eady_initial_v(state, theta0, rho0, v)

# Set initial u
u_exp = as_vector([u, v, 0.])
u0.project(u_exp)

# Pass these initial conditions to the state.initialise method
state.initialise([('u', u0),
                  ('rho', rho0),
                  ('theta', theta0)])

# Set the background profiles
state.set_reference_profiles([('rho', rho_b),
                              ('theta', theta_b)])

# Set up advection schemes
# we need a DG funciton space for the embedded DG advection scheme
ueqn = AdvectionEquation(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt, supg_params={"dg_direction": "horizontal"})

advected_fields = []
advected_fields.append(("u", SSPRK3(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))

# Set up linear solver
if hybrid:
    linear_solver_params = {'ksp_type': 'gmres',
                            'pc_type': 'gamg',
                            'ksp_monitor_true_residual': True,
                            'pc_gamg_sym_graph': True,
                            'mg_levels': {'ksp_type': 'chebyshev',
                                          'ksp_chebyshev_esteig': True,
                                          'ksp_max_it': 5,
                                          'pc_type': 'bjacobi',
                                          'sub_pc_type': 'ilu'}}
    linear_solver = HybridizedCompressibleSolver(state, solver_parameters=linear_solver_params,
                                                 overwrite_solver_parameters=True)

else:
    linear_solver_params = {
        'ksp_monitor_true_residual': True,
        'pc_type': 'fieldsplit',
        'pc_fieldsplit_type': 'schur',
        'ksp_type': 'gmres',
        'ksp_max_it': 100,
        'ksp_gmres_restart': 50,
        'pc_fieldsplit_schur_fact_type': 'FULL',
        'pc_fieldsplit_schur_precondition': 'selfp',
        'fieldsplit_0': {'ksp_type': 'preonly',
                         'pc_type': 'bjacobi',
                         'sub_pc_type': 'ilu'},
        'fieldsplit_1': {'ksp_type': 'preonly',
                         'pc_type': 'gamg',
                         'pc_gamg_sym_graph': True,
                         'mg_levels': {'ksp_type': 'chebyshev',
                                       'ksp_chebyshev_esteig': True,
                                       'ksp_max_it': 5,
                                       'pc_type': 'bjacobi',
                                       'sub_pc_type': 'ilu'}}
    }
    linear_solver = CompressibleSolver(state, solver_parameters=linear_solver_params,
                                       overwrite_solver_parameters=True)

# Set up forcing
forcing = CompressibleEadyForcing(state, euler_poincare=False)

# Build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, forcing)

PETSc.Sys.Print("Starting simulation...\n")
stepper.run(t=0, tmax=tmax)
<commit_msg>Functioning Eady slice model<commit_after>"""
Credit for demo setup: Gusto development team
"""
from gusto import *
from firedrake import as_vector, SpatialCoordinate,\
    PeriodicRectangleMesh, ExtrudedMesh, \
    exp, cos, sin, cosh, sinh, tanh, pi, Function, sqrt, \
    parameters
from firedrake.petsc import PETSc
from argparse import ArgumentParser
from hybridization import HybridizedCompressibleSolver
import sys

day = 24.*60.*60.
hour = 60.*60.

PETSc.Log.begin()
parser = ArgumentParser(description="""Euler-Boussinesq Eady slice model.""",
                        add_help=False)

parser.add_argument("--hybridization",
                    action="store_true",
                    help="Use a hybridized compressible solver.")

parser.add_argument("--test",
                    action="store_true",
                    help="Enable a quick test run.")

parser.add_argument("--profile",
                    action="store_true",
                    help="Turn on profiling for a 20 time-step run.")

parser.add_argument("--dt",
                    default=50.,
                    type=float,
                    action="store",
                    help="Time step size (seconds)")

parser.add_argument("--res",
                    default=30,
                    type=int,
                    action="store",
                    help="Resolution scaling parameter.")

parser.add_argument("--debug",
                    action="store_true",
                    help="Turn on KSP monitors")

parser.add_argument("--help",
                    action="store_true",
                    help="Show help.")

args, _ = parser.parse_known_args()

if args.help:
    help = parser.format_help()
    PETSc.Sys.Print("%s\n" % help)
    sys.exit(1)

dt = args.dt
res = args.res

if args.profile:
    # Ensures accurate timing of parallel loops
    parameters["pyop2_options"]["lazy_evaluation"] = False
    tmax = 20*dt

if args.test:
    tmax = dt

if not args.test and not args.profile:
    tmax = 30*day

hybrid = bool(args.hybridization)
PETSc.Sys.Print("""
Problem parameters:\n
Test case: Euler-Boussinesq Eady model.\n
Hybridized compressible solver: %s,\n
Profiling: %s,\n
Test run: %s.\n
""" % (hybrid, bool(args.profile), bool(args.test)))

PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" % (dt,
                                                                      tmax))

# Construct 1d periodic base mesh
columns = 2*res  # number of columns
L = 1000000.
m = PeriodicRectangleMesh(columns, 1, 2.*L, 1.e5, quadrilateral=True)

# Build 2D mesh by extruding the base mesh
nlayers = res    # horizontal layers
H = 10000.  # height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)

# Set up all the other things that state requires:
# Coriolis expression
f = 1.e-04
Omega = as_vector([0., 0., f*0.5])

# List of prognostic fieldnames:
# this is passed to state and used to construct a dictionary,
# state.field_dict so that we can access fields by name.
# u is the 3D velocity
# p is the pressure
# b is the buoyancy
fieldlist = ['u', 'rho', 'theta']

# Class containing timestepping parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
timestepping = TimesteppingParameters(dt=dt)

# Class containing output parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
if hybrid:
    dirname = "hybrid_compressible_eady"
else:
    dirname = "compressible_eady"

output = OutputParameters(dirname=dirname,
                          dumpfreq=int(2*hour/dt),
                          dumplist=['u', 'rho', 'theta'],
                          perturbation_fields=['rho', 'theta', 'ExnerPi'])

# Class containing physical parameters:
# all values not explicitly set here use the default values provided
# and documented in configuration.py
parameters = CompressibleEadyParameters(H=H, f=f)

# Class for diagnostics:
# fields passed to this class will have basic diagnostics computed
# (eg min, max, l2 norm) and these will be output as a json file
diagnostics = Diagnostics(*fieldlist)

# List of diagnostic fields, each defined in a class in diagnostics.py
diagnostic_fields = [CourantNumber(), VelocityY(),
                     ExnerPi(), ExnerPi(reference=True),
                     CompressibleKineticEnergy(),
                     CompressibleKineticEnergyY(),
                     CompressibleEadyPotentialEnergy(),
                     Sum("CompressibleKineticEnergy",
                         "CompressibleEadyPotentialEnergy"),
                     Difference("CompressibleKineticEnergy",
                                "CompressibleKineticEnergyY")]

# Setup state, passing in the mesh, information on the required finite element
# function spaces and the classes above
state = State(mesh, vertical_degree=1, horizontal_degree=1,
              family="RTCF",
              Coriolis=Omega,
              timestepping=timestepping,
              output=output,
              parameters=parameters,
              diagnostics=diagnostics,
              fieldlist=fieldlist,
              diagnostic_fields=diagnostic_fields)

# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")

# Spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()

# First setup the background buoyancy profile
# z.grad(bref) = N**2.
# The following is symbolic algebra, using the default buoyancy frequency
# from the parameters class.
x, y, z = SpatialCoordinate(mesh)
g = parameters.g
Nsq = parameters.Nsq
theta_surf = parameters.theta_surf

# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
theta_ref = theta_surf*exp(Nsq*(z-H/2)/g)
theta_b = Function(Vt).interpolate(theta_ref)

# Set theta_pert
def coth(x):
    return cosh(x)/sinh(x)

def Z(z):
    return Bu*((z/H)-0.5)

def n():
    return Bu**(-1)*sqrt((Bu*0.5-tanh(Bu*0.5))*(coth(Bu*0.5)-Bu*0.5))

a = -4.5
Bu = 0.5
theta_exp = a*theta_surf/g*sqrt(Nsq)*(-(1.-Bu*0.5*coth(Bu*0.5))*sinh(Z(z))*cos(pi*(x-L)/L)
                                      - n()*Bu*cosh(Z(z))*sin(pi*(x-L)/L))
theta_pert = Function(Vt).interpolate(theta_exp)

# Set theta0
theta0.interpolate(theta_b + theta_pert)

# Calculate hydrostatic Pi
PETSc.Sys.Print("Computing hydrostatic varaibles...\n")
params = {'pc_type': 'fieldsplit',
          'pc_fieldsplit_type': 'schur',
          'ksp_type': 'gmres',
          'ksp_monitor_true_residual': True,
          'ksp_rtol': 1.0e-8,
          'ksp_max_it': 1000,
          'ksp_gmres_restart': 50,
          'pc_fieldsplit_schur_fact_type': 'FULL',
          'pc_fieldsplit_schur_precondition': 'selfp',
          'fieldsplit_0': {'ksp_type': 'cg',
                           'ksp_rtol': 1.0e-8,
                           'pc_type': 'bjacobi',
                           'sub_pc_type': 'ilu'},
          'fieldsplit_1': {'ksp_type': 'cg',
                           'ksp_rtol': 1.0e-8,
                           'pc_type': 'gamg',
                           'pc_gamg_sym_graph': True,
                           'mg_levels': {'ksp_type': 'chebyshev',
                                         'ksp_chebyshev_esteig': True,
                                         'ksp_max_it': 5,
                                         'pc_type': 'bjacobi',
                                         'sub_pc_type': 'ilu'}}}

rho_b = Function(Vr)
compressible_hydrostatic_balance(state,
                                 theta_b,
                                 rho_b,
                                 params=params)
compressible_hydrostatic_balance(state,
                                 theta0,
                                 rho0,
                                 params=params)

# Set Pi0
Pi0 = calculate_Pi0(state, theta0, rho0)
state.parameters.Pi0 = Pi0

# Set x component of velocity
cp = state.parameters.cp
dthetady = state.parameters.dthetady
Pi = thermodynamics.pi(state.parameters, rho0, theta0)
u = cp*dthetady/f*(Pi-Pi0)

# Set y component of velocity
v = Function(Vr).assign(0.)
compressible_eady_initial_v(state, theta0, rho0, v)

# Set initial u
u_exp = as_vector([u, v, 0.])
u0.project(u_exp)

# Pass these initial conditions to the state.initialise method
state.initialise([('u', u0),
                  ('rho', rho0),
                  ('theta', theta0)])

# Set the background profiles
state.set_reference_profiles([('rho', rho_b),
                              ('theta', theta_b)])

# Set up advection schemes
# we need a DG funciton space for the embedded DG advection scheme
ueqn = AdvectionEquation(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt, supg_params={"dg_direction": "horizontal"})

advected_fields = []
advected_fields.append(("u", SSPRK3(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))

# Set up linear solver
if hybrid:
    linear_solver_params = {'ksp_type': 'gmres',
                            'pc_type': 'gamg',
                            'pc_gamg_sym_graph': True,
                            'mg_levels': {'ksp_type': 'chebyshev',
                                          'ksp_chebyshev_esteig': True,
                                          'ksp_max_it': 5,
                                          'pc_type': 'bjacobi',
                                          'sub_pc_type': 'ilu'}}
    if args.debug:
        linear_solver_params['ksp_monitor_true_residual'] = True

    linear_solver = HybridizedCompressibleSolver(state, solver_parameters=linear_solver_params,
                                                 overwrite_solver_parameters=True)

else:
    linear_solver_params = {'pc_type': 'fieldsplit',
                            'pc_fieldsplit_type': 'schur',
                            'ksp_type': 'gmres',
                            'ksp_max_it': 100,
                            'ksp_gmres_restart': 50,
                            'pc_fieldsplit_schur_fact_type': 'FULL',
                            'pc_fieldsplit_schur_precondition': 'selfp',
                            'fieldsplit_0': {'ksp_type': 'preonly',
                                             'pc_type': 'bjacobi',
                                             'sub_pc_type': 'ilu'},
                            'fieldsplit_1': {'ksp_type': 'preonly',
                                             'pc_type': 'gamg',
                                             'pc_gamg_sym_graph': True,
                                             'mg_levels': {'ksp_type': 'chebyshev',
                                                           'ksp_chebyshev_esteig': True,
                                                           'ksp_max_it': 5,
                                                           'pc_type': 'bjacobi',
                                                           'sub_pc_type': 'ilu'}}}
    if args.debug:
        linear_solver_params['ksp_monitor_true_residual'] = True

    linear_solver = CompressibleSolver(state, solver_parameters=linear_solver_params,
                                       overwrite_solver_parameters=True)

# Set up forcing
forcing = CompressibleEadyForcing(state, euler_poincare=False)

# Build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, forcing)

PETSc.Sys.Print("Starting simulation...\n")
stepper.run(t=0, tmax=tmax)
<|endoftext|>